"`):
+ 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
+
+ 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)
+
+
+__all__ = ["CodeLlamaTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/code_llama/tokenization_code_llama_fast.py b/lib/python3.12/site-packages/transformers/models/code_llama/tokenization_code_llama_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9d4c34b86ac577f2e6ec753450afdacc0704535
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/code_llama/tokenization_code_llama_fast.py
@@ -0,0 +1,381 @@
+# 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
+ [meta-llama/CodeLlama-7b-Instruct-hf](https://huggingface.co/meta-llama/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`):
+ Whether 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,)
+
+ 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
+
+
+__all__ = ["CodeLlamaTokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/__init__.py b/lib/python3.12/site-packages/transformers/models/dbrx/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..cce0f34c778d21a81d03940e7a7951707d898c86
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/dbrx/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_dbrx import *
+ from .modeling_dbrx import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f730a1f616c867188ff35d3205dc4588d7c7086c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7f614c28540af42e0e30ca2fd72e26d0351ca728
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..42cb3ecb72e72005ae7e610aeab011c385cacaaa
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/configuration_dbrx.py b/lib/python3.12/site-packages/transformers/models/dbrx/configuration_dbrx.py
new file mode 100644
index 0000000000000000000000000000000000000000..72df1fe335bab446fa1e88186d2ff14ae70f3351
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/dbrx/configuration_dbrx.py
@@ -0,0 +1,232 @@
+# 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 (`int`, *optional*, defaults to 1): For grouped_query_attention only, allow user to specify number of kv heads.
+ rope_theta (`float`, *optional*, defaults to 10000.0): The base frequency for rope.
+ """
+
+ base_config_key = "attn_config"
+
+ 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", "attn_implementation", "transformers_version", "_commit_hash", "torch_dtype"]:
+ if k in kwargs:
+ kwargs.pop(k)
+ if len(kwargs) != 0:
+ raise ValueError(f"Found unknown {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`, *optional*, defaults to 3584): The hidden size of the feedforward network.
+ moe_num_experts (`int`, *optional*, defaults to 4): The number of experts in the mixture of experts layer.
+ moe_top_k (`int`, *optional*, 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`, *optional*, 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.
+ """
+
+ base_config_key = "ffn_config"
+
+ 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", "attn_implementation", "transformers_version", "_commit_hash", "torch_dtype"]:
+ if k in kwargs:
+ kwargs.pop(k)
+ if len(kwargs) != 0:
+ raise ValueError(f"Found unknown {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"
+ sub_configs = {"attn_config": DbrxAttentionConfig, "ffn_config": DbrxFFNConfig}
+ 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
+ self.num_key_value_heads = self.attn_config.kv_n_heads
+
+ 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)
+
+
+__all__ = ["DbrxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/dbrx/modeling_dbrx.py b/lib/python3.12/site-packages/transformers/models/dbrx/modeling_dbrx.py
new file mode 100644
index 0000000000000000000000000000000000000000..258023c60b4d85913debc3df4cbf36582256fd55
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/dbrx/modeling_dbrx.py
@@ -0,0 +1,1395 @@
+# 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, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_torch_flex_attn_available,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.deprecation import deprecate_kwarg
+from .configuration_dbrx import DbrxConfig
+
+
+if is_torch_flex_attn_available():
+ from torch.nn.attention.flex_attention import BlockMask
+
+ from ...integrations.flex_attention import make_flex_block_causal_mask
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "DbrxConfig"
+
+
+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
+
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+ self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)
+
+ @torch.no_grad()
+ def forward(self, x, position_ids, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ self.inv_freq.to(x.device)
+ 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`, *optional*):
+ 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
+
+
+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)
+
+ 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, **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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ 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]]]:
+ if isinstance(past_key_value, 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"
+ )
+ 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)
+
+ 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 = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ position_ids=position_ids,
+ dropout=dropout_rate,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ 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
+
+
+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)
+
+ 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()
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ is_causal = True if causal_mask is None and q_len > 1 else False
+
+ 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,
+ is_causal=is_causal,
+ )
+
+ 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):
+ # (This cause torch.compile to fail with `torch._dynamo.exc.Unsupported: dynamic shape operator: aten.nonzero.default`)
+ # (set torch._dynamo.config.capture_dynamic_output_shape_ops = True may help but not tested)
+ 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: Optional[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
+ _supports_quantized_cache = True
+ _supports_static_cache = False # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
+
+ 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.fill_(1.0)
+ 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)
+
+
+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, see our
+ [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);
+ - 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,
+ **kwargs, # NOOP kwargs, for now
+ ) -> 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 must specify exactly one of input_ids or inputs_embeds")
+
+ 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)
+
+ # kept for BC (non `Cache` `past_key_values` inputs)
+ return_legacy_cache = False
+ if use_cache and not isinstance(past_key_values, Cache):
+ return_legacy_cache = True
+ if past_key_values is None:
+ past_key_values = DynamicCache()
+ else:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ logger.warning_once(
+ "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
+ "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
+ "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
+ )
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ 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_key_values, output_attentions
+ )
+
+ # 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 = next_decoder_cache if use_cache else None
+ if return_legacy_cache:
+ next_cache = next_cache.to_legacy_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,
+ )
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask
+ def _update_causal_mask(
+ self,
+ attention_mask: torch.Tensor,
+ input_tensor: torch.Tensor,
+ cache_position: torch.Tensor,
+ past_key_values: Cache,
+ output_attentions: bool = False,
+ ):
+ if self.config._attn_implementation == "flash_attention_2":
+ if attention_mask is not None and (attention_mask == 0.0).any():
+ return attention_mask
+ return None
+ if self.config._attn_implementation == "flex_attention":
+ if isinstance(attention_mask, torch.Tensor):
+ attention_mask = make_flex_block_causal_mask(attention_mask)
+ if isinstance(attention_mask, BlockMask):
+ return attention_mask
+
+ # 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. This feature is not compatible with static cache, as SDPA will fail
+ # to infer the attention mask.
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ using_static_cache = isinstance(past_key_values, StaticCache)
+
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
+ attention_mask,
+ inputs_embeds=input_tensor,
+ past_key_values_length=past_seen_tokens,
+ is_training=self.training,
+ ):
+ return None
+
+ dtype, device = input_tensor.dtype, input_tensor.device
+ sequence_length = input_tensor.shape[1]
+ if using_static_cache:
+ target_length = past_key_values.get_max_cache_shape()
+ else:
+ target_length = (
+ attention_mask.shape[-1]
+ if isinstance(attention_mask, torch.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask,
+ sequence_length=sequence_length,
+ target_length=target_length,
+ dtype=dtype,
+ device=device,
+ cache_position=cache_position,
+ batch_size=input_tensor.shape[0],
+ )
+
+ if (
+ self.config._attn_implementation == "sdpa"
+ and attention_mask is not None
+ and attention_mask.device.type in ["cuda", "xpu"]
+ and not output_attentions
+ ):
+ # 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
+ min_dtype = torch.finfo(dtype).min
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+ return causal_mask
+
+ @staticmethod
+ # Copied from transformers.models.llama.modeling_llama.LlamaPreTrainedModel._prepare_4d_causal_attention_mask_with_cache_position
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ **kwargs,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+ `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache,
+ to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ min_dtype = torch.finfo(dtype).min
+ 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(batch_size, 1, -1, -1)
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+ causal_mask.device
+ )
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+
+@add_start_docstrings("The DBRX Model transformer for causal language modeling.", DBRX_START_DOCSTRING)
+class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin):
+ 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
+
+ @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+ @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,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs,
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+ 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+ 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]
+ # No upscaling to float was ever done for Dbrx
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ 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,
+ )
+
+
+__all__ = ["DbrxForCausalLM", "DbrxModel", "DbrxPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/__init__.py b/lib/python3.12/site-packages/transformers/models/deit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..98236a86d7a1e8b4ff16b53fb3ff37befbf1d7ac
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/__init__.py
@@ -0,0 +1,31 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_deit import *
+ from .feature_extraction_deit import *
+ from .image_processing_deit import *
+ from .image_processing_deit_fast import *
+ from .modeling_deit import *
+ from .modeling_tf_deit import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/deit/configuration_deit.py b/lib/python3.12/site-packages/transformers/models/deit/configuration_deit.py
new file mode 100644
index 0000000000000000000000000000000000000000..63d772067c5fd3e13ab0051b377751a873c7f03b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/configuration_deit.py
@@ -0,0 +1,152 @@
+# 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__)
+
+
+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.
+ pooler_output_size (`int`, *optional*):
+ Dimensionality of the pooler layer. If None, defaults to `hidden_size`.
+ pooler_act (`str`, *optional*, defaults to `"tanh"`):
+ The activation function to be used by the pooler. Keys of ACT2FN are supported for Flax and
+ Pytorch, and elements of https://www.tensorflow.org/api_docs/python/tf/keras/activations are
+ supported for Tensorflow.
+
+ 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,
+ pooler_output_size=None,
+ pooler_act="tanh",
+ **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
+ self.pooler_output_size = pooler_output_size if pooler_output_size else hidden_size
+ self.pooler_act = pooler_act
+
+
+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
+
+
+__all__ = ["DeiTConfig", "DeiTOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/feature_extraction_deit.py b/lib/python3.12/site-packages/transformers/models/deit/feature_extraction_deit.py
new file mode 100644
index 0000000000000000000000000000000000000000..813c115fce553dfbf6d22efa41b123c50e6113c7
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/feature_extraction_deit.py
@@ -0,0 +1,36 @@
+# 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)
+
+
+__all__ = ["DeiTFeatureExtractor"]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit.py b/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1eceebcb5c0c302e2d6c8bcea95d9c15f33dfeb
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit.py
@@ -0,0 +1,299 @@
+# 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_preprocess_arguments,
+)
+from ...utils import TensorType, filter_out_non_signature_kwargs, 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
+
+ # 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,
+ )
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample=None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: 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,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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)
+
+ 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 do_rescale and is_scaled_image(images[0]):
+ 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])
+
+ all_images = []
+ for image in images:
+ 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
+ )
+
+ all_images.append(image)
+ images = [
+ to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+ for image in all_images
+ ]
+
+ data = {"pixel_values": images}
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["DeiTImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit_fast.py b/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..28cd6539df795d230abca5c70d4d8e07922b869c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/image_processing_deit_fast.py
@@ -0,0 +1,44 @@
+# 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.
+"""Fast Image processor class for DeiT."""
+
+from ...image_processing_utils_fast import BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, BaseImageProcessorFast
+from ...image_utils import (
+ IMAGENET_STANDARD_MEAN,
+ IMAGENET_STANDARD_STD,
+ PILImageResampling,
+)
+from ...utils import add_start_docstrings
+
+
+@add_start_docstrings(
+ "Constructs a fast DeiT image processor.",
+ BASE_IMAGE_PROCESSOR_FAST_DOCSTRING,
+)
+class DeiTImageProcessorFast(BaseImageProcessorFast):
+ # To be checked against the slow image processor
+ # None values left after checking can be removed
+ resample = PILImageResampling.BICUBIC
+ image_mean = IMAGENET_STANDARD_MEAN
+ image_std = IMAGENET_STANDARD_STD
+ size = {"height": 256, "width": 256}
+ crop_size = {"height": 224, "width": 224}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+
+
+__all__ = ["DeiTImageProcessorFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/modeling_deit.py b/lib/python3.12/site-packages/transformers/models/deit/modeling_deit.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb88bca353fc551afe959231476e64c97cd30206
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/modeling_deit.py
@@ -0,0 +1,999 @@
+# 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
+from dataclasses import dataclass
+from typing import Callable, 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 ALL_ATTENTION_FUNCTIONS, 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,
+ torch_int,
+)
+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"
+
+
+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)
+ self.patch_size = config.patch_size
+
+ 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. This method is also adapted to support torch.jit tracing and 2 class embeddings.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 2
+ num_positions = self.position_embeddings.shape[1] - 2
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_and_dist_pos_embed = self.position_embeddings[:, :2]
+ patch_pos_embed = self.position_embeddings[:, 2:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_and_dist_pos_embed, 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:
+ _, _, height, width = pixel_values.shape
+ 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)
+ position_embedding = self.position_embeddings
+
+ if interpolate_pos_encoding:
+ position_embedding = self.interpolate_pos_encoding(embeddings, height, width)
+
+ embeddings = embeddings + position_embedding
+ 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."
+ )
+ x = self.projection(pixel_values).flatten(2).transpose(1, 2)
+ return x
+
+
+# Copied from transformers.models.vit.modeling_vit.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+
+ # Normalize the attention scores to probabilities.
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ # Mask heads if we want to
+ if attention_mask is not None:
+ attn_weights = attn_weights * attention_mask
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# 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.config = config
+ 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.dropout_prob = config.attention_probs_dropout_prob
+ self.scaling = self.attention_head_size**-0.5
+ self.is_causal = False
+
+ 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)
+
+ 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]]:
+ 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(self.query(hidden_states))
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ if self.config._attn_implementation == "sdpa" and output_attentions:
+ logger.warning_once(
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ else:
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ context_layer, attention_probs = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ head_mask,
+ is_causal=self.is_causal,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.dropout_prob,
+ )
+
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.reshape(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,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"]
+ _supports_sdpa = True
+ _supports_flash_attn_2 = 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, DeiTEmbeddings):
+ module.cls_token.data.zero_()
+ module.position_embeddings.data.zero_()
+ module.distillation_token.data.zero_()
+ if module.mask_token is not None:
+ module.mask_token.data.zero_()
+
+
+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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
+ Whether to interpolate the pre-trained position encodings.
+"""
+
+
+@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,
+ interpolate_pos_encoding: bool = False,
+ ) -> 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, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ 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.pooler_output_size)
+ self.activation = ACT2FN[config.pooler_act]
+
+ 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,
+ interpolate_pos_encoding: bool = False,
+ ) -> 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ 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,
+ interpolate_pos_encoding: bool = False,
+ ) -> 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ 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: Optional[torch.FloatTensor] = None
+ cls_logits: Optional[torch.FloatTensor] = None
+ distillation_logits: Optional[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,
+ interpolate_pos_encoding: bool = False,
+ ) -> 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ 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,
+ )
+
+
+__all__ = [
+ "DeiTForImageClassification",
+ "DeiTForImageClassificationWithTeacher",
+ "DeiTForMaskedImageModeling",
+ "DeiTModel",
+ "DeiTPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/deit/modeling_tf_deit.py b/lib/python3.12/site-packages/transformers/models/deit/modeling_tf_deit.py
new file mode 100644
index 0000000000000000000000000000000000000000..49c95268035a98c55e45071080ef748bce1d17ba
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/deit/modeling_tf_deit.py
@@ -0,0 +1,1233 @@
+# 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"
+
+
+@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: Optional[tf.Tensor] = None
+ cls_logits: Optional[tf.Tensor] = None
+ distillation_logits: Optional[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 interpolate_pos_encoding(self, embeddings: tf.Tensor, height: int, width: int) -> tf.Tensor:
+ num_patches = embeddings.shape[1] - 2
+ num_positions = self.position_embeddings.shape[1] - 2
+
+ if num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, 0, :]
+ dist_pos_embed = self.position_embeddings[:, 1, :]
+ patch_pos_embed = self.position_embeddings[:, 2:, :]
+ dim = embeddings.shape[-1]
+ h0 = height // self.config.patch_size
+ w0 = 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
+ h0, w0 = h0 + 0.1, w0 + 0.1
+ patch_pos_embed = tf.reshape(
+ patch_pos_embed, (1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
+ )
+ patch_pos_embed = tf.image.resize(patch_pos_embed, size=(int(h0), int(w0)), method="bicubic")
+ patch_pos_embed = tf.transpose(patch_pos_embed, perm=[0, 2, 3, 1])
+ patch_pos_embed = tf.reshape(patch_pos_embed, (1, -1, dim))
+
+ return tf.concat(
+ [tf.expand_dims(class_pos_embed, axis=0), tf.expand_dims(dist_pos_embed, axis=0), patch_pos_embed], axis=1
+ )
+
+ def call(
+ self,
+ pixel_values: tf.Tensor,
+ bool_masked_pos: tf.Tensor | None = None,
+ training: bool = False,
+ interpolate_pos_encoding: bool = False,
+ ) -> tf.Tensor:
+ _, height, width, _ = pixel_values.shape
+
+ 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)
+ position_embedding = self.position_embeddings
+ if interpolate_pos_encoding:
+ position_embedding = self.interpolate_pos_encoding(embeddings, height, width)
+
+ embeddings = embeddings + position_embedding
+ 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."
+ )
+
+ 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,
+ interpolate_pos_encoding: bool = False,
+ 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ 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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
+ Whether to interpolate the pre-trained position encodings.
+ 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,
+ interpolate_pos_encoding: bool = False,
+ 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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.pooler_output_size,
+ kernel_initializer=get_initializer(config.initializer_range),
+ activation=config.pooler_act,
+ 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,
+ interpolate_pos_encoding: bool = False,
+ 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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,
+ interpolate_pos_encoding: bool = False,
+ 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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,
+ interpolate_pos_encoding: bool = False,
+ 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,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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])
+
+
+__all__ = [
+ "TFDeiTForImageClassification",
+ "TFDeiTForImageClassificationWithTeacher",
+ "TFDeiTForMaskedImageModeling",
+ "TFDeiTModel",
+ "TFDeiTPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__init__.py b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c786feb9213fdd31640c0fdeaead5164026ad37a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__init__.py
@@ -0,0 +1,29 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_encoder_decoder import *
+ from .modeling_encoder_decoder import *
+ from .modeling_flax_encoder_decoder import *
+ from .modeling_tf_encoder_decoder import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a4ab3b646d3a0733a49abd2d6e716543bbfe558c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2d37a5b992da896471b30ea53ae0e9e339466a09
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..970f8defe2cf3263fb8d9927c549182388c85c73
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..acb19ae3b9f8b50f53320e71cfd6ef5cfc80918e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c64446c631bf3df8c2c5d05bbb8ceed7b2112d64
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/configuration_encoder_decoder.py b/lib/python3.12/site-packages/transformers/models/encoder_decoder/configuration_encoder_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b5c62363f6ad9c3ce72c61bd597dac3f9078c0a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/encoder_decoder/configuration_encoder_decoder.py
@@ -0,0 +1,111 @@
+# 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
+from ..auto import AutoConfig
+
+
+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"
+ sub_configs = {"encoder": AutoConfig, "decoder": AutoConfig}
+ is_composition = True
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ if "encoder" not in kwargs or "decoder" not in kwargs:
+ raise ValueError(
+ f"A configuraton of type {self.model_type} cannot be instantiated because "
+ f"both `encoder` and `decoder` sub-configurations were not passed, only {kwargs}"
+ )
+ 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")
+
+ 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)
+
+
+__all__ = ["EncoderDecoderConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_encoder_decoder.py b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_encoder_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..415fd058e45d8fe6e44ca30ddfe178ffc717cf19
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_encoder_decoder.py
@@ -0,0 +1,689 @@
+# 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 ...generation import GenerationMixin
+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, GenerationMixin):
+ 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
+ _supports_param_buffer_assignment = False
+ _supports_flash_attn_2 = True
+ _supports_sdpa = 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
+ # update `_attn_implementation` because the attn is set in a deepcopied config within PreTrainedModel
+ self.config.encoder._attn_implementation = self.encoder.config._attn_implementation
+ self.config.decoder._attn_implementation = self.decoder.config._attn_implementation
+ 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):
+ self.encoder.tie_weights()
+ self.decoder.tie_weights()
+ # 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 _init_weights(self, module):
+ if module in self.encoder.modules():
+ self.encoder._init_weights(module)
+ elif module in self.decoder.modules():
+ self.decoder._init_weights(module)
+
+ 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
+
+ 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: Optional[str] = None,
+ decoder_pretrained_model_name_or_path: Optional[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 "num_items_in_batch" in kwargs_encoder:
+ kwargs_decoder["num_items_in_batch"] = kwargs_encoder.pop("num_items_in_batch", None)
+
+ 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 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)
+
+
+__all__ = ["EncoderDecoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..bdc589484cda1083f42ce96e67b12771f84b6af9
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py
@@ -0,0 +1,901 @@
+# 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
+
+
+__all__ = ["FlaxEncoderDecoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5abafc361b6b4e71e365d7b7f3caa3368f82d92
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py
@@ -0,0 +1,665 @@
+# 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: Optional[str] = None,
+ decoder_pretrained_model_name_or_path: Optional[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)
+
+
+__all__ = ["TFEncoderDecoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__init__.py b/lib/python3.12/site-packages/transformers/models/esm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8eac54d6ddcbdae2b8ca3771ae5540522f6f29da
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_esm import *
+ from .modeling_esm import *
+ from .modeling_esmfold import *
+ from .modeling_tf_esm import *
+ from .tokenization_esm import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..de507d9778c41b7001934e621f130314b4df14fa
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__pycache__/configuration_esm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/configuration_esm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bfe1a632f1901722a83953fcce5589baaff774c1
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/configuration_esm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_esm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_esm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c594265f351474b9dcda912bf2d47bfd912908bf
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_esm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b960d0c8b25100b670d3d44c833439118c0fa529
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/__pycache__/tokenization_esm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/tokenization_esm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3e418309d4194ad86cd53fbe0f014acb4a890254
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/__pycache__/tokenization_esm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/configuration_esm.py b/lib/python3.12/site-packages/transformers/models/esm/configuration_esm.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac56bc8d783a5895496520fc62c7bada7fa5f61a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/configuration_esm.py
@@ -0,0 +1,365 @@
+# 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
+
+
+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(vocab_size=33)
+
+ >>> # 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: Optional[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",
+ ".",
+ "-",
+ "",
+ "",
+ )
+
+
+__all__ = ["EsmConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/esm/modeling_esm.py b/lib/python3.12/site-packages/transformers/models/esm/modeling_esm.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f90d8d052f9bc9a3f0c5cdeb74aad823b44bf21
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/modeling_esm.py
@@ -0,0 +1,1273 @@
+# 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"
+
+
+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 with BertLMPredictionHead->EsmLMHead
+ 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)
+ elif isinstance(module, EsmLMHead):
+ module.bias.data.zero_()
+
+
+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
+
+
+__all__ = [
+ "EsmForMaskedLM",
+ "EsmForSequenceClassification",
+ "EsmForTokenClassification",
+ "EsmModel",
+ "EsmPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/esm/modeling_esmfold.py b/lib/python3.12/site-packages/transformers/models/esm/modeling_esmfold.py
new file mode 100644
index 0000000000000000000000000000000000000000..645c9d16a5c5dcd012cf3d91ca784267ede9c654
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/modeling_esmfold.py
@@ -0,0 +1,2325 @@
+# 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: Optional[torch.FloatTensor] = None
+ sidechain_frames: Optional[torch.FloatTensor] = None
+ unnormalized_angles: Optional[torch.FloatTensor] = None
+ angles: Optional[torch.FloatTensor] = None
+ positions: Optional[torch.FloatTensor] = None
+ states: Optional[torch.FloatTensor] = None
+ s_s: Optional[torch.FloatTensor] = None
+ s_z: Optional[torch.FloatTensor] = None
+ distogram_logits: Optional[torch.FloatTensor] = None
+ lm_logits: Optional[torch.FloatTensor] = None
+ aatype: Optional[torch.FloatTensor] = None
+ atom14_atom_exists: Optional[torch.FloatTensor] = None
+ residx_atom14_to_atom37: Optional[torch.FloatTensor] = None
+ residx_atom37_to_atom14: Optional[torch.FloatTensor] = None
+ atom37_atom_exists: Optional[torch.FloatTensor] = None
+ residue_index: Optional[torch.FloatTensor] = None
+ lddt_head: Optional[torch.FloatTensor] = None
+ plddt: Optional[torch.FloatTensor] = None
+ ptm_logits: Optional[torch.FloatTensor] = None
+ ptm: Optional[torch.FloatTensor] = None
+ aligned_confidence_probs: Optional[torch.FloatTensor] = None
+ predicted_aligned_error: Optional[torch.FloatTensor] = None
+ max_predicted_aligned_error: Optional[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 sequences (.. 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 (dtype=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)
+
+
+__all__ = ["EsmForProteinFolding", "EsmFoldPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/esm/modeling_tf_esm.py b/lib/python3.12/site-packages/transformers/models/esm/modeling_tf_esm.py
new file mode 100644
index 0000000000000000000000000000000000000000..71698486dab0adfece1feb14ce91b4500e24382c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/modeling_tf_esm.py
@@ -0,0 +1,1575 @@
+# 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
+
+
+__all__ = [
+ "TFEsmForMaskedLM",
+ "TFEsmForSequenceClassification",
+ "TFEsmForTokenClassification",
+ "TFEsmModel",
+ "TFEsmPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__init__.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..02a8c149ae320dd9b045edc5df31760a4eebefd9
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__init__.py
@@ -0,0 +1,8 @@
+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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8fcddfec4a0f263acc4899fcb1f583bbf2331b17
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ce2f6290ceb5ead4d79446dd679ddddd1e00f670
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8fa3df18b9c6de681474dc00d7a7367fd40e46e0
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..246786c961d3c5a0b7760253ce0e79be416fb390
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5f870bd2f3a32eebc509e41bdba4aef21dae8a59
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..403415cabac5b628e764c4c1f873631f4662bc13
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1724b07c8f34e0e35bfdb9e3fd1828a992e78535
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5c2a0345ee9f1181e242b3aeb233fd03b1063182
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f60cd4c47dbab7ef07129fe1debe0fd323e59949
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/chunk_utils.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/chunk_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..51ff6b74d6c3f59ac23b68e2be0999b5f8dbb4ef
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/chunk_utils.py
@@ -0,0 +1,397 @@
+# 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 TypeError("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 TypeError("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) is 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/data_transforms.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/data_transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d4c17589ae66df2a8fd0ccfe8d6e335004eed9a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/data_transforms.py
@@ -0,0 +1,93 @@
+# 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/feats.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/feats.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac7b90dfe79b24b852cb26fca998bda831f36a6f
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/feats.py
@@ -0,0 +1,253 @@
+# 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/loss.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c442786dc82ba2ebe243923509ed76a40de2a01
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/loss.py
@@ -0,0 +1,105 @@
+# 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/protein.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/protein.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae9d8c13277bd8c5e7dd152f50d549b6f7286af3
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/protein.py
@@ -0,0 +1,330 @@
+# 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/residue_constants.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/residue_constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..b05a603fb29f51e2b870c73fa28899cd91936bdb
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/residue_constants.py
@@ -0,0 +1,981 @@
+# 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 TypeError("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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d0f2f69b350ce7589a1c82e896cb08ec5808f41
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py
@@ -0,0 +1,1242 @@
+# 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) is not 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) is not 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 homogeneous transformation tensor.
+
+ Returns:
+ A [*, 4, 4] homogeneous 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 homogeneous transformation tensor.
+
+ Args:
+ t: [*, 4, 4] homogeneous 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/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/tensor_utils.py b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/tensor_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..efe72e4905b81faee5fe44131f1ba1b75856bd1c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/tensor_utils.py
@@ -0,0 +1,140 @@
+# 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 TypeError("Not supported")
+
+
+tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)
diff --git a/lib/python3.12/site-packages/transformers/models/esm/tokenization_esm.py b/lib/python3.12/site-packages/transformers/models/esm/tokenization_esm.py
new file mode 100644
index 0000000000000000000000000000000000000000..4bc433e350e13c33fc779092baf52197d8aa5e0d
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/esm/tokenization_esm.py
@@ -0,0 +1,147 @@
+# 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)
+
+
+__all__ = ["EsmTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..00b6ccc53fc0efb0fc88c2f95586276cd40010fe
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py
@@ -0,0 +1,32 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_got_ocr2 import *
+ from .image_processing_got_ocr2 import *
+ from .image_processing_got_ocr2_fast import *
+ from .modeling_got_ocr2 import *
+ from .processing_got_ocr2 import *
+
+
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..19683baac85bed19698afd89aa344b44e76a0bdc
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/configuration_got_ocr2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/configuration_got_ocr2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ffc775aaf6153fcd18a4d7a5ce6259a15114df5e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/configuration_got_ocr2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7215a41cf49ad7f63ac9d3b98ea79b9cba8071ed
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2bc30f6a21c0ea03625ddbc34b7b65f7dc693573
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/image_processing_got_ocr2_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modeling_got_ocr2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modeling_got_ocr2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b26902e3c04b7450aae7dc544ab0f6b415eb8777
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modeling_got_ocr2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modular_got_ocr2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modular_got_ocr2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e0c986da5398b4852404eb7015536d93c6cd8650
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/modular_got_ocr2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/processing_got_ocr2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/processing_got_ocr2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..657ae908f9ff654819c2a4360ba4c5fc17beb9f8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/got_ocr2/__pycache__/processing_got_ocr2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb9a1fb68889ca1013c46da8fc04a7d301e7e059
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py
@@ -0,0 +1,208 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_got_ocr2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 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 ...configuration_utils import PretrainedConfig
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+class GotOcr2VisionConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`GotOcr2VisionModel`]. It is used to instantiate a GOT_OCR2
+ vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
+ defaults will yield a similar configuration to that of the SAM ViT-h
+ [facebook/sam-vit-huge](https://huggingface.co/facebook/sam-vit-huge) 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.
+ output_channels (`int`, *optional*, defaults to 256):
+ Dimensionality of the output channels in the Patch 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.
+ num_channels (`int`, *optional*, defaults to 3):
+ Number of channels in the input image.
+ image_size (`int`, *optional*, defaults to 1024):
+ Expected resolution. Target size of the resized input image.
+ patch_size (`int`, *optional*, defaults to 16):
+ Size of the patches to be extracted from the input image.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string)
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ 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 1e-10):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ qkv_bias (`bool`, *optional*, defaults to `True`):
+ Whether to add a bias to query, key, value projections.
+ use_abs_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use absolute position embedding.
+ use_rel_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use relative position embedding.
+ window_size (`int`, *optional*, defaults to 14):
+ Window size for relative position.
+ global_attn_indexes (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`):
+ The indexes of the global attention layers.
+ mlp_dim (`int`, *optional*, defaults to 3072):
+ The dimensionality of the MLP layer in the Transformer encoder.
+ """
+
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ hidden_size=768,
+ output_channels=256,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ num_channels=3,
+ image_size=1024,
+ patch_size=16,
+ hidden_act="gelu",
+ layer_norm_eps=1e-06,
+ attention_dropout=0.0,
+ initializer_range=1e-10,
+ qkv_bias=True,
+ use_abs_pos=True,
+ use_rel_pos=True,
+ window_size=14,
+ global_attn_indexes=[2, 5, 8, 11],
+ mlp_dim=3072,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.hidden_size = hidden_size
+ self.output_channels = output_channels
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_channels = num_channels
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.attention_dropout = attention_dropout
+ self.initializer_range = initializer_range
+ self.qkv_bias = qkv_bias
+ self.use_abs_pos = use_abs_pos
+ self.use_rel_pos = use_rel_pos
+ self.window_size = window_size
+ self.global_attn_indexes = global_attn_indexes
+ self.mlp_dim = mlp_dim
+
+
+class GotOcr2Config(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`GotOcr2ForConditionalGeneration`]. It is used to instantiate a
+ GotOcr2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of GOT-OCR-2.0.
+
+ e.g [stepfun-ai/GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf)
+
+ 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 (`Union[AutoConfig, dict]`, *optional*, defaults to `CLIPVisionConfig`):
+ The config object or dictionary of the vision backbone.
+ text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`):
+ The config object or dictionary of the text backbone.
+ image_token_index (`int`, *optional*, defaults to 151859):
+ The image token index to encode the image prompt.
+ image_seq_length (`int`, *optional*, defaults to 576):
+ Sequence length of one image embedding.
+ pad_token_id (`int`, *optional*, defaults to -1):
+ Padding token id.
+
+ ```python
+ >>> from transformers import GotOcr2ForConditionalGeneration, GotOcr2Config
+
+ >>> # Initializing a GotOcr2 style configuration
+ >>> configuration = GotOcr2Config()
+
+ >>> # Initializing a model from the Qwen2-VL-7B style configuration
+ >>> model = GotOcr2ForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "got_ocr2"
+ sub_configs = {"text_config": AutoConfig, "vision_config": GotOcr2VisionConfig}
+
+ def __init__(
+ self,
+ vision_config=None,
+ text_config=None,
+ image_token_index=151859,
+ image_seq_length=576,
+ pad_token_id=-1,
+ **kwargs,
+ ):
+ self.image_token_index = image_token_index
+ self.image_seq_length = image_seq_length
+ self.pad_token_id = pad_token_id
+
+ if vision_config is None:
+ self.vision_config = GotOcr2VisionConfig()
+ elif isinstance(vision_config, dict):
+ self.vision_config = GotOcr2VisionConfig(**vision_config)
+ elif isinstance(vision_config, GotOcr2VisionConfig):
+ self.vision_config = vision_config
+
+ if isinstance(text_config, dict):
+ text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "qwen2"
+ text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+ elif text_config is None:
+ text_config = CONFIG_MAPPING["qwen2"](
+ vocab_size=151860,
+ hidden_size=1024,
+ intermediate_size=2816,
+ num_hidden_layers=24,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=True,
+ rope_theta=1000000.0,
+ rope_scaling=None,
+ use_sliding_window=False,
+ sliding_window=4096,
+ max_window_layers=21,
+ attention_dropout=0.0,
+ )
+
+ self.text_config = text_config
+
+ super().__init__(**kwargs)
+
+
+__all__ = ["GotOcr2VisionConfig", "GotOcr2Config"]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..875c0742b96d7c0cea0049b932c68c4254fe71b2
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py
@@ -0,0 +1,495 @@
+# coding=utf-8
+# Copyright 2024 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 Got-OCR-2."""
+
+from functools import lru_cache
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from ...image_transforms import (
+ convert_to_rgb,
+ 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_flat_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
+
+
+if is_vision_available():
+ import PIL
+
+
+logger = logging.get_logger(__name__)
+
+
+# Similar to image_processing_mllama.get_all_supported_aspect_ratios
+@lru_cache(maxsize=10)
+def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> List[Tuple[int, int]]:
+ """
+ Computes all allowed aspect ratios for a given minimum and maximum number of input tiles.
+
+ This function calculates all possible arrangements of tiles that can be formed
+ within the constraint of the minimum and maximum number of tiles. Each arrangement is
+ represented by its aspect ratio (width/height) and the corresponding tile configuration.
+
+ Args:
+ min_image_tiles (`int`):
+ The minimum number of tiles allowed.
+ max_image_tiles (`int`):
+ The maximum number of tiles allowed.
+
+ Returns:
+ `List[Tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height)
+ configuration in terms of number of tiles.
+
+ Example:
+ >>> get_all_supported_aspect_ratios(1, 4)
+ [(1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (2, 2), (4, 1)]
+
+ """
+ aspect_ratios = []
+ for width in range(1, max_image_tiles + 1):
+ for height in range(1, max_image_tiles + 1):
+ if width * height <= max_image_tiles and width * height >= min_image_tiles:
+ aspect_ratios.append((width, height))
+
+ aspect_ratios = sorted(aspect_ratios, key=lambda x: x[0] * x[1])
+
+ return aspect_ratios
+
+
+@lru_cache(maxsize=100)
+def get_optimal_tiled_canvas(
+ original_image_size: Tuple[int, int],
+ target_tile_size: Tuple[int, int],
+ min_image_tiles: int,
+ max_image_tiles: int,
+) -> Tuple[int, int]:
+ """
+ Given a minimum and maximum number of tiles, find the canvas with the closest aspect ratio to the
+ original image aspect ratio.
+ In case of tie-breaking condition when two canvases have the same aspect ratio difference, we favor the canvas with
+ more tiles, until the area covered by the tiles is more than twice the target area, in order to avoid unnecessarily
+ excessive tiling.
+ """
+ possible_tile_arrangements = get_all_supported_aspect_ratios(min_image_tiles, max_image_tiles)
+
+ original_height, original_width = original_image_size
+ target_tile_height, target_tile_width = target_tile_size
+ aspect_ratio = original_width / original_height
+ area = original_width * original_height
+
+ # find the grid with the best aspect ratio
+ best_ratio_diff = float("inf")
+ best_grid = (1, 1)
+ for grid in possible_tile_arrangements:
+ grid_aspect_ratio = grid[0] / grid[1]
+ ratio_diff = abs(aspect_ratio - grid_aspect_ratio)
+ if ratio_diff < best_ratio_diff:
+ best_ratio_diff = ratio_diff
+ best_grid = grid
+ elif ratio_diff == best_ratio_diff:
+ # if the aspect ratio difference is the same, we favor the grid with more patches
+ # until the area covered by the patches is more than twice the original image area
+ if area > 0.5 * target_tile_height * target_tile_width * grid[0] * grid[1]:
+ best_grid = grid
+
+ return best_grid
+
+
+class GotOcr2ImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a GOT_OCR2 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 the `preprocess` method.
+ size (`dict`, *optional*, defaults to `{"height": 384, "width": 384}`):
+ Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
+ method.
+ crop_to_patches (`bool`, *optional*, defaults to `False`):
+ Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the
+ `preprocess` method.
+ min_patches (`int`, *optional*, defaults to 1):
+ The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method.
+ max_patches (`int`, *optional*, defaults to 12):
+ The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method.
+ resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+ Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be
+ overridden by the `resample` 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.
+ rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+ Scale factor to use if rescaling the image. Only has an effect if `do_rescale` is set to `True`. Can be
+ overridden by the `rescale_factor` 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. 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. 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.
+ Can be overridden by the `image_std` parameter in the `preprocess` method.
+ do_convert_rgb (`bool`, *optional*, defaults to `True`):
+ Whether to convert the image to RGB.
+ """
+
+ model_input_names = ["pixel_values"]
+
+ def __init__(
+ self,
+ do_resize: bool = True,
+ size: Dict[str, int] = None,
+ crop_to_patches: bool = False,
+ min_patches: int = 1,
+ max_patches: int = 12,
+ resample: PILImageResampling = PILImageResampling.BICUBIC,
+ 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_convert_rgb: bool = True,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ size = size if size is not None else {"height": 384, "width": 384}
+ size = get_size_dict(size, default_to_square=True)
+
+ self.do_resize = do_resize
+ self.size = size
+ self.crop_to_patches = crop_to_patches
+ self.min_patches = min_patches
+ self.max_patches = max_patches
+ 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 OPENAI_CLIP_MEAN
+ self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+ self.do_convert_rgb = do_convert_rgb
+
+ 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,
+ )
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Optional[Dict[str, int]] = None,
+ crop_to_patches: Optional[bool] = None,
+ min_patches: Optional[int] = None,
+ max_patches: Optional[int] = None,
+ resample: PILImageResampling = 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,
+ do_convert_rgb: Optional[bool] = None,
+ data_format: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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`):
+ Controls the size of the image after `resize`. The shortest edge of the image is resized to
+ `size["shortest_edge"]` whilst preserving the aspect ratio. If the longest edge of this resized image
+ is > `int(size["shortest_edge"] * (1333 / 800))`, then the image is resized again to make the longest
+ edge equal to `int(size["shortest_edge"] * (1333 / 800))`.
+ crop_to_patches (`bool`, *optional*, defaults to `self.crop_to_patches`):
+ Whether to crop the image to patches.
+ min_patches (`int`, *optional*, defaults to `self.min_patches`):
+ The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`.
+ max_patches (`int`, *optional*, defaults to `self.max_patches`):
+ The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`.
+ resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+ Resampling filter to use if resizing the image. 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 to normalize the image by if `do_normalize` is set to `True`.
+ image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+ Image standard deviation to normalize the image by if `do_normalize` is set to `True`.
+ do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+ Whether to convert the image to RGB.
+ 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_to_patches = crop_to_patches if crop_to_patches is not None else self.crop_to_patches
+ min_patches = min_patches if min_patches is not None else self.min_patches
+ max_patches = max_patches if max_patches is not None else self.max_patches
+ 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_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+ size = size if size is not None else self.size
+ size = get_size_dict(size, default_to_square=False)
+ images = make_flat_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,
+ )
+ # PIL RGBA images are converted to RGB
+ if do_convert_rgb:
+ images = [convert_to_rgb(image) for image in images]
+
+ # All transformations expect numpy arrays.
+ images = [to_numpy_array(image) for image in images]
+
+ if do_rescale and is_scaled_image(images[0]):
+ 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 crop_to_patches and max_patches > 1:
+ images = [
+ self.crop_image_to_patches(
+ image,
+ min_patches=min_patches,
+ max_patches=max_patches,
+ patch_size=size,
+ data_format=input_data_format,
+ )
+ for image in images
+ ]
+ num_patches = np.array([len(image) for image in images])
+ images = [image for images_list in images for image in images_list]
+ else:
+ num_patches = np.array([1] * len(images))
+
+ for i, image in enumerate(images):
+ if do_resize:
+ images[i] = self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
+
+ if do_rescale:
+ images[i] = self.rescale(image=images[i], scale=rescale_factor, input_data_format=input_data_format)
+
+ if do_normalize:
+ images[i] = self.normalize(
+ image=images[i],
+ mean=image_mean,
+ std=image_std,
+ input_data_format=input_data_format,
+ )
+
+ images[i] = to_channel_dimension_format(images[i], data_format, input_channel_dim=input_data_format)
+
+ encoded_outputs = BatchFeature(
+ data={"pixel_values": images, "num_patches": num_patches}, tensor_type=return_tensors
+ )
+
+ return encoded_outputs
+
+ def crop_image_to_patches(
+ self,
+ images: np.ndarray,
+ min_patches: int,
+ max_patches: int,
+ use_thumbnail: bool = True,
+ patch_size: Union[Tuple, int, dict] = None,
+ data_format: ChannelDimension = None,
+ ):
+ """
+ Crop the image to patches and return a list of cropped images.
+ The number of patches and their grid arrangement are determined by the original image size,
+ the target patch size and the minimum and maximum number of patches.
+ The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio.
+
+ Args:
+ images (`np.ndarray`):
+ The image to be cropped.
+ min_patches (`int`):
+ The minimum number of patches to be extracted from the image.
+ max_patches (`int`):
+ The maximum number of patches to be extracted from the image.
+ use_thumbnail (`bool`, *optional*, defaults to `True`):
+ Whether to add a thumbnail image to the list of cropped patches.
+ patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*):
+ The size of the output patches.
+ data_format (`ChannelDimension`, *optional*):
+ The format of the image data. If `None`, the format is inferred from the input image.
+
+ Returns:
+ List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images.
+ """
+ if data_format is None:
+ data_format = infer_channel_dimension_format(images)
+ images = to_channel_dimension_format(images, ChannelDimension.FIRST, data_format)
+ patch_size_height, patch_size_width = patch_size["height"], patch_size["width"]
+ original_height, original_width = images.shape[-2:]
+ # find the closest aspect ratio to the target
+ num_columns, num_rows = get_optimal_tiled_canvas(
+ (original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches
+ )
+
+ # calculate the target width and height
+ target_width = patch_size_width * num_columns
+ target_height = patch_size_height * num_rows
+ num_blocks = num_columns * num_rows
+
+ # resize the image so that each patch is of patch_size
+ resized_image = self.resize(
+ images,
+ {"height": target_height, "width": target_width},
+ data_format=ChannelDimension.FIRST,
+ input_data_format=ChannelDimension.FIRST,
+ )
+ # split the image into patches
+ processed_images = []
+ for i in range(num_blocks):
+ column = i % num_columns
+ row = i // num_columns
+ box = (
+ column * patch_size_width,
+ row * patch_size_height,
+ (column + 1) * patch_size_width,
+ (row + 1) * patch_size_height,
+ )
+ # split the image
+ patch_image = resized_image[..., box[1] : box[3], box[0] : box[2]]
+ patch_image = to_channel_dimension_format(patch_image, data_format, ChannelDimension.FIRST)
+ processed_images.append(patch_image)
+
+ if use_thumbnail and len(processed_images) != 1:
+ thumbnail_img = self.resize(
+ images, patch_size, data_format=data_format, input_data_format=ChannelDimension.FIRST
+ )
+ processed_images.append(thumbnail_img)
+
+ return processed_images
+
+
+__all__ = ["GotOcr2ImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..8498e378030868a30dd7dac537999d3aa8a39a63
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2_fast.py
@@ -0,0 +1,249 @@
+# coding=utf-8
+# Copyright 2025 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.
+"""Fast Image processor class for Got-OCR-2."""
+
+from typing import List, Optional, Tuple, Union
+
+from ...image_processing_utils import BatchFeature
+from ...image_processing_utils_fast import (
+ BASE_IMAGE_PROCESSOR_FAST_DOCSTRING,
+ BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS,
+ BaseImageProcessorFast,
+ DefaultFastImageProcessorKwargs,
+ group_images_by_shape,
+ reorder_images,
+)
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import Unpack
+from ...utils import (
+ TensorType,
+ add_start_docstrings,
+ is_torch_available,
+ is_torchvision_available,
+ is_torchvision_v2_available,
+)
+from .image_processing_got_ocr2 import get_optimal_tiled_canvas
+
+
+if is_torch_available():
+ import torch
+
+if is_torchvision_available():
+ if is_torchvision_v2_available():
+ from torchvision.transforms.v2 import functional as F
+ else:
+ from torchvision.transforms import functional as F
+
+
+class GotOcr2ImageProcessorKwargs(DefaultFastImageProcessorKwargs):
+ crop_to_patches: Optional[bool]
+ min_patches: Optional[int]
+ max_patches: Optional[int]
+
+
+@add_start_docstrings(
+ "Constructs a fast GotOcr2 image processor.",
+ BASE_IMAGE_PROCESSOR_FAST_DOCSTRING,
+ """
+ crop_to_patches (`bool`, *optional*, defaults to `False`):
+ Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the
+ `preprocess` method.
+ min_patches (`int`, *optional*, defaults to 1):
+ The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method.
+ max_patches (`int`, *optional*, defaults to 12):
+ The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method.
+ """,
+)
+class GotOcr2ImageProcessorFast(BaseImageProcessorFast):
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"height": 384, "width": 384}
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+ crop_to_patches = False
+ min_patches = 1
+ max_patches = 12
+ valid_kwargs = GotOcr2ImageProcessorKwargs
+
+ def __init__(self, **kwargs: Unpack[valid_kwargs]):
+ super().__init__(**kwargs)
+
+ @add_start_docstrings(
+ BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS,
+ """
+ crop_to_patches (`bool`, *optional*, defaults to `False`):
+ Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the
+ `preprocess` method.
+ min_patches (`int`, *optional*, defaults to 1):
+ The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method.
+ max_patches (`int`, *optional*, defaults to 12):
+ The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method.
+ """,
+ )
+ def preprocess(self, images: ImageInput, **kwargs: Unpack[valid_kwargs]) -> BatchFeature:
+ return super().preprocess(images, **kwargs)
+
+ def crop_image_to_patches(
+ self,
+ images: "torch.Tensor",
+ min_patches: int,
+ max_patches: int,
+ use_thumbnail: bool = True,
+ patch_size: Union[Tuple, int, dict] = None,
+ interpolation: Optional["F.InterpolationMode"] = None,
+ ):
+ """
+ Crop the images to patches and return a list of cropped images.
+ The number of patches and their grid arrangement are determined by the original image size,
+ the target patch size and the minimum and maximum number of patches.
+ The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio.
+
+ Args:
+ images (`torch.Tensor`):
+ The images to be cropped.
+ min_patches (`int`):
+ The minimum number of patches to be extracted from the image.
+ max_patches (`int`):
+ The maximum number of patches to be extracted from the image.
+ use_thumbnail (`bool`, *optional*, defaults to `True`):
+ Whether to add a thumbnail image to the list of cropped patches.
+ patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*):
+ The size of the output patches.
+ The format of the image data. If `None`, the format is inferred from the input image.
+
+ Returns:
+ List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images.
+ """
+ patch_size_height, patch_size_width = patch_size.height, patch_size.width
+ original_height, original_width = images.shape[-2:]
+ # find the closest aspect ratio to the target
+ num_columns, num_rows = get_optimal_tiled_canvas(
+ (original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches
+ )
+
+ # calculate the target width and height
+ target_width = patch_size_width * num_columns
+ target_height = patch_size_height * num_rows
+ num_blocks = num_columns * num_rows
+
+ # resize the image so that each patch is of patch_size
+ resized_image = self.resize(
+ images, SizeDict(height=target_height, width=target_width), interpolation=interpolation
+ )
+ # split the image into patches
+ processed_images = []
+ for i in range(num_blocks):
+ column = i % num_columns
+ row = i // num_columns
+ box = (
+ column * patch_size_width,
+ row * patch_size_height,
+ (column + 1) * patch_size_width,
+ (row + 1) * patch_size_height,
+ )
+ # split the image
+ patch_image = resized_image[..., box[1] : box[3], box[0] : box[2]]
+ processed_images.append(patch_image)
+
+ if use_thumbnail and len(processed_images) != 1:
+ thumbnail_img = self.resize(images, patch_size, interpolation=interpolation)
+ processed_images.append(thumbnail_img)
+
+ processed_images = torch.stack(processed_images, dim=0).transpose(0, 1).contiguous()
+
+ return processed_images
+
+ def _preprocess(
+ self,
+ images: List["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ crop_to_patches: bool,
+ min_patches: int,
+ max_patches: int,
+ interpolation: Optional["F.InterpolationMode"],
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: Optional[Union[float, List[float]]],
+ image_std: Optional[Union[float, List[float]]],
+ return_tensors: Optional[Union[str, TensorType]],
+ ) -> BatchFeature:
+ if crop_to_patches:
+ grouped_images, grouped_images_index = group_images_by_shape(images)
+ processed_images_grouped = {}
+ num_patches = {}
+ for shape, stacked_images in grouped_images.items():
+ stacked_images = self.crop_image_to_patches(
+ stacked_images,
+ min_patches,
+ max_patches,
+ patch_size=size,
+ interpolation=interpolation,
+ )
+ processed_images_grouped[shape] = stacked_images
+ num_patches[shape] = [stacked_images.shape[1]] * stacked_images.shape[0]
+ images = reorder_images(processed_images_grouped, grouped_images_index)
+ images = [image for images_list in images for image in images_list]
+ num_patches = reorder_images(num_patches, grouped_images_index)
+ else:
+ num_patches = [1] * len(images)
+
+ # Group images by size for batched resizing
+ grouped_images, grouped_images_index = group_images_by_shape(images)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(image=stacked_images, size=size, interpolation=interpolation)
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ # Group images by size for further processing
+ # Needed in case do_resize is False, or resize returns images with different sizes
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_center_crop:
+ stacked_images = self.center_crop(stacked_images, crop_size)
+ # Fused rescale and normalize
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+ processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images
+
+ return BatchFeature(
+ data={"pixel_values": processed_images, "num_patches": num_patches}, tensor_type=return_tensors
+ )
+
+
+__all__ = ["GotOcr2ImageProcessorFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..c474ef36900fee4c3af6b45a59e86c0949d60618
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py
@@ -0,0 +1,903 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_got_ocr2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 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 collections
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_outputs import ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ can_return_tuple,
+ replace_return_docstrings,
+)
+from ..auto import AutoModelForCausalLM
+from .configuration_got_ocr2 import GotOcr2Config, GotOcr2VisionConfig
+
+
+_CONFIG_FOR_DOC = "GotOcr2Config"
+
+
+class GotOcr2MLPBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim)
+ self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size)
+ self.act = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.lin1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.lin2(hidden_states)
+ return hidden_states
+
+
+class GotOcr2VisionAttention(nn.Module):
+ """Multi-head Attention block with relative position embeddings."""
+
+ def __init__(self, config, window_size):
+ super().__init__()
+ input_size = (
+ (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ if window_size == 0
+ else (window_size, window_size)
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ head_dim = config.hidden_size // config.num_attention_heads
+ self.scale = head_dim**-0.5
+ self.dropout = config.attention_dropout
+
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size)
+
+ self.use_rel_pos = config.use_rel_pos
+ if self.use_rel_pos:
+ if input_size is None:
+ raise ValueError("Input size must be provided if using relative positional encoding.")
+
+ # initialize relative positional embeddings
+ self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
+ self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
+
+ def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
+ """
+ Get relative positional embeddings according to the relative positions of
+ query and key sizes.
+
+ Args:
+ q_size (int):
+ size of the query.
+ k_size (int):
+ size of key k.
+ rel_pos (`torch.Tensor`):
+ relative position embeddings (L, channel).
+
+ Returns:
+ Extracted positional embeddings according to relative positions.
+ """
+ max_rel_dist = int(2 * max(q_size, k_size) - 1)
+ # Interpolate rel pos.
+ rel_pos_resized = F.interpolate(
+ rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
+ size=max_rel_dist,
+ mode="linear",
+ )
+ rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+
+ # Scale the coords with short length if shapes for q and k are different.
+ q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+ k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+ relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+ return rel_pos_resized[relative_coords.long()]
+
+ def get_decomposed_rel_pos(
+ self,
+ query: torch.Tensor,
+ rel_pos_h: torch.Tensor,
+ rel_pos_w: torch.Tensor,
+ q_size: Tuple[int, int],
+ k_size: Tuple[int, int],
+ ) -> torch.Tensor:
+ """
+ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+ https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py
+
+ Args:
+ query (`torch.Tensor`):
+ query q in the attention layer with shape (batch_size, query_height * query_width, channel).
+ rel_pos_h (`torch.Tensor`):
+ relative position embeddings (Lh, channel) for height axis.
+ rel_pos_w (`torch.Tensor`):
+ relative position embeddings (Lw, channel) for width axis.
+ q_size (tuple):
+ spatial sequence size of query q with (query_height, query_width).
+ k_size (tuple):
+ spatial sequence size of key k with (key_height, key_width).
+
+ Returns:
+ decomposed_rel_pos (`torch.Tensor`):
+ decomposed relative position embeddings.
+ """
+ query_height, query_width = q_size
+ key_height, key_width = k_size
+ relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h)
+ relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w)
+
+ batch_size, _, dim = query.shape
+ reshaped_query = query.reshape(batch_size, query_height, query_width, dim)
+ rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height)
+ rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width)
+
+ decomposed_rel_pos = rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
+
+ return decomposed_rel_pos
+
+ def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch.Tensor:
+ batch_size, height, width, _ = hidden_states.shape
+ # qkv with shape (3, batch_size, nHead, height * width, channel)
+ qkv = (
+ self.qkv(hidden_states)
+ .reshape(batch_size, height * width, 3, self.num_attention_heads, -1)
+ .permute(2, 0, 3, 1, 4)
+ )
+ # q, k, v with shape (batch_size * nHead, height * width, channel)
+ query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0)
+
+ attn_weights = (query * self.scale) @ key.transpose(-2, -1)
+
+ if self.use_rel_pos:
+ decomposed_rel_pos = self.get_decomposed_rel_pos(
+ query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width)
+ )
+ decomposed_rel_pos = decomposed_rel_pos.reshape_as(attn_weights)
+ attn_weights = attn_weights + decomposed_rel_pos
+
+ attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype)
+
+ attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+ attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1)
+ attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1)
+
+ attn_output = self.proj(attn_output)
+
+ if output_attentions:
+ outputs = (attn_output, attn_weights)
+ else:
+ outputs = (attn_output, None)
+
+ return outputs
+
+
+class GotOcr2VisionLayer(nn.Module):
+ def __init__(self, config, window_size):
+ super().__init__()
+ self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn = GotOcr2VisionAttention(config, window_size)
+ self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp = GotOcr2MLPBlock(config)
+ self.window_size = window_size
+
+ def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
+ """
+ Args:
+ Partition into non-overlapping windows with padding if needed.
+ hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window
+ size.
+
+ Returns:
+ windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel].
+ (pad_height, pad_width): padded height and width before partition
+ """
+ batch_size, height, width, channel = hidden_states.shape
+
+ pad_h = (window_size - height % window_size) % window_size
+ pad_w = (window_size - width % window_size) % window_size
+ hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h))
+ pad_height, pad_width = height + pad_h, width + pad_w
+
+ hidden_states = hidden_states.reshape(
+ batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel
+ )
+ windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel)
+ return windows, (pad_height, pad_width)
+
+ def window_unpartition(
+ self, windows: torch.Tensor, window_size: int, padding_shape: Tuple[int, int], original_shape: Tuple[int, int]
+ ) -> torch.Tensor:
+ """
+ Args:
+ Window unpartition into original sequences and removing padding.
+ hidden_states (tensor):
+ input tokens with [batch_size * num_windows, window_size, window_size, channel].
+ window_size (int):
+ window size.
+ padding_shape (Tuple):
+ padded height and width (pad_height, pad_width).
+ original_shape (Tuple): original height and width (height, width) before padding.
+
+ Returns:
+ hidden_states: unpartitioned sequences with [batch_size, height, width, channel].
+ """
+ pad_height, pad_width = padding_shape
+ height, width = original_shape
+ batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size)
+ hidden_states = windows.reshape(
+ batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1
+ )
+ hidden_states = (
+ hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1)
+ )
+
+ hidden_states = hidden_states[:, :height, :width, :].contiguous()
+ return hidden_states
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: Optional[bool] = False,
+ ) -> Tuple[torch.FloatTensor]:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ # Window partition
+ if self.window_size > 0:
+ height, width = hidden_states.shape[1], hidden_states.shape[2]
+ hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size)
+
+ hidden_states, attn_weights = self.attn(
+ hidden_states=hidden_states,
+ output_attentions=output_attentions,
+ )
+ # Reverse window partition
+ if self.window_size > 0:
+ hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width))
+
+ hidden_states = residual + hidden_states
+ layernorm_output = self.layer_norm2(hidden_states)
+ hidden_states = hidden_states + self.mlp(layernorm_output)
+
+ outputs = (hidden_states,)
+ if output_attentions:
+ outputs += (attn_weights,)
+
+ return outputs
+
+
+@dataclass
+class GotOcr2VisionEncoderOutput(ModelOutput):
+ """
+ Base class for got_ocr2 vision model's outputs that also contains image embeddings obtained by applying the projection
+ layer to the pooler_output.
+
+ 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: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+class GotOcr2PatchEmbeddings(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):
+ 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]})."
+ )
+ embeddings = self.projection(pixel_values).permute(0, 2, 3, 1)
+ return embeddings
+
+
+class GotOcr2LayerNorm(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 GotOcr2VisionNeck(nn.Module):
+ def __init__(self, config: GotOcr2VisionConfig):
+ super().__init__()
+ self.config = config
+
+ self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False)
+ self.layer_norm1 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")
+ self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False)
+ self.layer_norm2 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")
+
+ def forward(self, hidden_states):
+ hidden_states = hidden_states.permute(0, 3, 1, 2)
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.layer_norm1(hidden_states)
+
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layer_norm2(hidden_states)
+ return hidden_states
+
+
+class GotOcr2VisionEncoder(nn.Module):
+ def __init__(self, config: GotOcr2VisionConfig):
+ super().__init__()
+ self.config = config
+ self.image_size = config.image_size
+
+ self.patch_embed = GotOcr2PatchEmbeddings(config)
+
+ self.pos_embed = None
+ if config.use_abs_pos:
+ # Initialize absolute positional embedding with pretrain image size.
+ self.pos_embed = nn.Parameter(
+ torch.zeros(
+ 1,
+ config.image_size // config.patch_size,
+ config.image_size // config.patch_size,
+ config.hidden_size,
+ )
+ )
+
+ self.layers = nn.ModuleList()
+ for i in range(config.num_hidden_layers):
+ layer = GotOcr2VisionLayer(
+ config,
+ window_size=config.window_size if i not in config.global_attn_indexes else 0,
+ )
+ self.layers.append(layer)
+
+ self.neck = GotOcr2VisionNeck(config)
+
+ self.gradient_checkpointing = False
+
+ def get_input_embeddings(self):
+ return self.patch_embed
+
+ @can_return_tuple
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ ) -> GotOcr2VisionEncoderOutput:
+ 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 pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.patch_embed(pixel_values)
+ if self.pos_embed is not None:
+ hidden_states = hidden_states + self.pos_embed
+
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ for i, layer_module in enumerate(self.layers):
+ 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, output_attentions=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,)
+
+ hidden_states = self.neck(hidden_states)
+
+ return GotOcr2VisionEncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+class GotOcr2MultiModalProjector(nn.Module):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__()
+ vision_output_channels = config.vision_config.output_channels
+ language_hidden_size = config.text_config.hidden_size
+ self.conv_upsampler1 = nn.Conv2d(
+ vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.conv_upsampler2 = nn.Conv2d(
+ vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size)
+
+ def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.conv_upsampler1(vision_embeddings)
+ hidden_state = self.conv_upsampler2(hidden_state)
+ hidden_state = hidden_state.flatten(2).permute(0, 2, 1)
+ hidden_state = self.multimodal_projector(hidden_state)
+ return hidden_state
+
+
+@dataclass
+class GotOcr2CausalLMOutputWithPast(ModelOutput):
+ """
+ Base class for GotOcr2 causal language model (or autoregressive) outputs.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token 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).
+ 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)`)
+
+ 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.
+ 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_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size (batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+ image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+GOT_OCR2_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 ([`GotOcr2Config`] or [`GotOcr2VisionConfig`]):
+ 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 LLaMA Model outputting raw hidden-states without any specific head on top.",
+ GOT_OCR2_START_DOCSTRING,
+)
+class GotOcr2PreTrainedModel(PreTrainedModel):
+ config_class = GotOcr2Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GotOcr2VisionAttention"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_cache_class = True
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_quantized_cache = True
+ _supports_static_cache = True
+
+ def _init_weights(self, module):
+ # important: this ported version of GotOcr2 isn't meant for training from scratch - only
+ # inference and fine-tuning - so the proper init weights code has been removed - the original codebase
+ # https://github.com/haotian-liu/GotOcr2/tree/main/got_ocr2 should serve for that purpose
+ std = (
+ self.config.initializer_range
+ if hasattr(self.config, "initializer_range")
+ else self.config.text_config.initializer_range
+ )
+
+ if hasattr(module, "class_embedding"):
+ module.class_embedding.data.normal_(mean=0.0, std=std)
+
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ 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_()
+
+
+GOT_OCR2_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)
+ pixel_values (`torch.FloatTensor` of shape `(seq_length, num_channels * image_size * image_size)):
+ The tensors corresponding to the input images. Pixel values can be obtained using
+ [`AutoImageProcessor`]. See [`GotOcr2ImageProcessor.__call__`] for details. [`GotOcr2Processor`] uses
+ [`GotOcr2ImageProcessor`] for processing images.
+ 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.
+ 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 GOT_OCR2 model which consists of a vision backbone and a language model.""",
+ GOT_OCR2_START_DOCSTRING,
+)
+class GotOcr2ForConditionalGeneration(GotOcr2PreTrainedModel, GenerationMixin):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__(config)
+ self.vision_tower = GotOcr2VisionEncoder(config.vision_config)
+
+ self.multi_modal_projector = GotOcr2MultiModalProjector(config)
+ self.vocab_size = config.text_config.vocab_size
+ self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+ if self.language_model._tied_weights_keys is not None:
+ self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys]
+
+ self.pad_token_id = config.pad_token_id
+
+ 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 get_output_embeddings(self):
+ return self.language_model.get_output_embeddings()
+
+ def set_output_embeddings(self, new_embeddings):
+ self.language_model.set_output_embeddings(new_embeddings)
+
+ def set_decoder(self, decoder):
+ self.language_model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.language_model.get_decoder()
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ ):
+ """
+ Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+ Args:
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+ Returns:
+ image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+ """
+ image_outputs = self.vision_tower(pixel_values).last_hidden_state
+ return self.multi_modal_projector(image_outputs)
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(GOT_OCR2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=GotOcr2CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = 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,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ ) -> GotOcr2CausalLMOutputWithPast:
+ 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer
+
+ >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda")
+ >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf")
+
+ >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda")
+
+ >>> # Generate
+ >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+ >>> generate_ids = model.generate(
+ ... **inputs,
+ ... do_sample=False,
+ ... tokenizer = processor.tokenizer,
+ ... stop_strings='<|im_end|>',
+ ... streamer=streamer,
+ ... max_new_tokens=4096,
+ ... )
+ "You should keep in mind what features from the module should be used, especially
+ when you're planning to sell a template."
+ ```"""
+
+ 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 (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if pixel_values is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype))
+ n_image_tokens = (input_ids == self.config.image_token_index).sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ if n_image_tokens != n_image_features:
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
+ special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+ image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs: CausalLMOutputWithPast = self.language_model(
+ 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,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ )
+
+ logits = outputs.logits
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ if attention_mask is not None:
+ # we use the input attention mask to shift the logits and labels, because it is 2D.
+ # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+ shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+ shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+ shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+ else:
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = nn.CrossEntropyLoss()
+ loss = loss_fct(
+ shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+ )
+
+ return GotOcr2CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ inputs_embeds=None,
+ pixel_values=None,
+ attention_mask=None,
+ cache_position=None,
+ logits_to_keep=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = self.language_model.prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ **kwargs,
+ )
+
+ if cache_position[0] == 0:
+ # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+ # Otherwise we need pixel values to be passed to model
+ model_inputs["pixel_values"] = pixel_values
+
+ return model_inputs
+
+
+__all__ = ["GotOcr2PreTrainedModel", "GotOcr2ForConditionalGeneration"]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/modular_got_ocr2.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/modular_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..4b7e7f1adb5e1393d4edf978e0d68d2bbadf7fe2
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/modular_got_ocr2.py
@@ -0,0 +1,522 @@
+# coding=utf-8
+# Copyright 2024 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 List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.llava.modeling_llava import (
+ LlavaCausalLMOutputWithPast,
+ LlavaForConditionalGeneration,
+ LlavaPreTrainedModel,
+)
+from transformers.models.sam.modeling_sam import SamMLPBlock, SamVisionAttention, SamVisionEncoder, SamVisionLayer
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import (
+ add_start_docstrings_to_model_forward,
+ can_return_tuple,
+ is_vision_available,
+ logging,
+ replace_return_docstrings,
+)
+from ..auto import CONFIG_MAPPING, AutoConfig, AutoModelForCausalLM
+
+
+if is_vision_available():
+ pass
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "GotOcr2Config"
+
+
+class GotOcr2VisionConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`GotOcr2VisionModel`]. It is used to instantiate a GOT_OCR2
+ vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
+ defaults will yield a similar configuration to that of the SAM ViT-h
+ [facebook/sam-vit-huge](https://huggingface.co/facebook/sam-vit-huge) 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.
+ output_channels (`int`, *optional*, defaults to 256):
+ Dimensionality of the output channels in the Patch 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.
+ num_channels (`int`, *optional*, defaults to 3):
+ Number of channels in the input image.
+ image_size (`int`, *optional*, defaults to 1024):
+ Expected resolution. Target size of the resized input image.
+ patch_size (`int`, *optional*, defaults to 16):
+ Size of the patches to be extracted from the input image.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string)
+ layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+ 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 1e-10):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ qkv_bias (`bool`, *optional*, defaults to `True`):
+ Whether to add a bias to query, key, value projections.
+ use_abs_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use absolute position embedding.
+ use_rel_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use relative position embedding.
+ window_size (`int`, *optional*, defaults to 14):
+ Window size for relative position.
+ global_attn_indexes (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`):
+ The indexes of the global attention layers.
+ mlp_dim (`int`, *optional*, defaults to 3072):
+ The dimensionality of the MLP layer in the Transformer encoder.
+ """
+
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ hidden_size=768,
+ output_channels=256,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ num_channels=3,
+ image_size=1024,
+ patch_size=16,
+ hidden_act="gelu",
+ layer_norm_eps=1e-06,
+ attention_dropout=0.0,
+ initializer_range=1e-10,
+ qkv_bias=True,
+ use_abs_pos=True,
+ use_rel_pos=True,
+ window_size=14,
+ global_attn_indexes=[2, 5, 8, 11],
+ mlp_dim=3072,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.hidden_size = hidden_size
+ self.output_channels = output_channels
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_channels = num_channels
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.attention_dropout = attention_dropout
+ self.initializer_range = initializer_range
+ self.qkv_bias = qkv_bias
+ self.use_abs_pos = use_abs_pos
+ self.use_rel_pos = use_rel_pos
+ self.window_size = window_size
+ self.global_attn_indexes = global_attn_indexes
+ self.mlp_dim = mlp_dim
+
+
+class GotOcr2Config(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`GotOcr2ForConditionalGeneration`]. It is used to instantiate a
+ GotOcr2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of GOT-OCR-2.0.
+
+ e.g [stepfun-ai/GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf)
+
+ 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 (`Union[AutoConfig, dict]`, *optional*, defaults to `CLIPVisionConfig`):
+ The config object or dictionary of the vision backbone.
+ text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`):
+ The config object or dictionary of the text backbone.
+ image_token_index (`int`, *optional*, defaults to 151859):
+ The image token index to encode the image prompt.
+ image_seq_length (`int`, *optional*, defaults to 576):
+ Sequence length of one image embedding.
+ pad_token_id (`int`, *optional*, defaults to -1):
+ Padding token id.
+
+ ```python
+ >>> from transformers import GotOcr2ForConditionalGeneration, GotOcr2Config
+
+ >>> # Initializing a GotOcr2 style configuration
+ >>> configuration = GotOcr2Config()
+
+ >>> # Initializing a model from the Qwen2-VL-7B style configuration
+ >>> model = GotOcr2ForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "got_ocr2"
+ sub_configs = {"text_config": AutoConfig, "vision_config": GotOcr2VisionConfig}
+
+ def __init__(
+ self,
+ vision_config=None,
+ text_config=None,
+ image_token_index=151859,
+ image_seq_length=576,
+ pad_token_id=-1,
+ **kwargs,
+ ):
+ self.image_token_index = image_token_index
+ self.image_seq_length = image_seq_length
+ self.pad_token_id = pad_token_id
+
+ if vision_config is None:
+ self.vision_config = GotOcr2VisionConfig()
+ elif isinstance(vision_config, dict):
+ self.vision_config = GotOcr2VisionConfig(**vision_config)
+ elif isinstance(vision_config, GotOcr2VisionConfig):
+ self.vision_config = vision_config
+
+ if isinstance(text_config, dict):
+ text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "qwen2"
+ text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+ elif text_config is None:
+ text_config = CONFIG_MAPPING["qwen2"](
+ vocab_size=151860,
+ hidden_size=1024,
+ intermediate_size=2816,
+ num_hidden_layers=24,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=True,
+ rope_theta=1000000.0,
+ rope_scaling=None,
+ use_sliding_window=False,
+ sliding_window=4096,
+ max_window_layers=21,
+ attention_dropout=0.0,
+ )
+
+ self.text_config = text_config
+
+ super().__init__(**kwargs)
+
+
+class GotOcr2MLPBlock(SamMLPBlock):
+ pass
+
+
+class GotOcr2VisionAttention(SamVisionAttention):
+ pass
+
+
+class GotOcr2VisionLayer(SamVisionLayer):
+ def __init__(self, config, window_size):
+ super().__init__()
+ self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn = GotOcr2VisionAttention(config, window_size)
+ self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp = GotOcr2MLPBlock(config)
+ self.window_size = window_size
+
+
+class GotOcr2VisionEncoder(SamVisionEncoder):
+ pass
+
+
+class GotOcr2MultiModalProjector(nn.Module):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__()
+ vision_output_channels = config.vision_config.output_channels
+ language_hidden_size = config.text_config.hidden_size
+ self.conv_upsampler1 = nn.Conv2d(
+ vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.conv_upsampler2 = nn.Conv2d(
+ vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size)
+
+ def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.conv_upsampler1(vision_embeddings)
+ hidden_state = self.conv_upsampler2(hidden_state)
+ hidden_state = hidden_state.flatten(2).permute(0, 2, 1)
+ hidden_state = self.multimodal_projector(hidden_state)
+ return hidden_state
+
+
+class GotOcr2CausalLMOutputWithPast(LlavaCausalLMOutputWithPast):
+ pass
+
+
+class GotOcr2PreTrainedModel(LlavaPreTrainedModel):
+ pass
+
+
+GOT_OCR2_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)
+ pixel_values (`torch.FloatTensor` of shape `(seq_length, num_channels * image_size * image_size)):
+ The tensors corresponding to the input images. Pixel values can be obtained using
+ [`AutoImageProcessor`]. See [`GotOcr2ImageProcessor.__call__`] for details. [`GotOcr2Processor`] uses
+ [`GotOcr2ImageProcessor`] for processing images.
+ 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.
+ 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.
+"""
+
+
+class GotOcr2ForConditionalGeneration(LlavaForConditionalGeneration):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__(config)
+ self.vision_tower = GotOcr2VisionEncoder(config.vision_config)
+
+ self.multi_modal_projector = GotOcr2MultiModalProjector(config)
+ self.vocab_size = config.text_config.vocab_size
+ self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+ if self.language_model._tied_weights_keys is not None:
+ self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys]
+
+ self.pad_token_id = config.pad_token_id
+
+ self.post_init()
+
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ ):
+ """
+ Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+ Args:
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+ Returns:
+ image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+ """
+ image_outputs = self.vision_tower(pixel_values).last_hidden_state
+ return self.multi_modal_projector(image_outputs)
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(GOT_OCR2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=GotOcr2CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = 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,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ ) -> GotOcr2CausalLMOutputWithPast:
+ 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer
+
+ >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda")
+ >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf")
+
+ >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda")
+
+ >>> # Generate
+ >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+ >>> generate_ids = model.generate(
+ ... **inputs,
+ ... do_sample=False,
+ ... tokenizer = processor.tokenizer,
+ ... stop_strings='<|im_end|>',
+ ... streamer=streamer,
+ ... max_new_tokens=4096,
+ ... )
+ "You should keep in mind what features from the module should be used, especially
+ when you're planning to sell a template."
+ ```"""
+
+ 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 (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if pixel_values is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype))
+ n_image_tokens = (input_ids == self.config.image_token_index).sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ if n_image_tokens != n_image_features:
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
+ special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+ image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs: CausalLMOutputWithPast = self.language_model(
+ 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,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ )
+
+ logits = outputs.logits
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ if attention_mask is not None:
+ # we use the input attention mask to shift the logits and labels, because it is 2D.
+ # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+ shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+ shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+ shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+ else:
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = nn.CrossEntropyLoss()
+ loss = loss_fct(
+ shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+ )
+
+ return GotOcr2CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+__all__ = [
+ "GotOcr2VisionConfig",
+ "GotOcr2Config",
+ "GotOcr2PreTrainedModel",
+ "GotOcr2ForConditionalGeneration",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/got_ocr2/processing_got_ocr2.py b/lib/python3.12/site-packages/transformers/models/got_ocr2/processing_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..398ec36c9e75a9ee101bf3aeacf26257aebccc9e
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/got_ocr2/processing_got_ocr2.py
@@ -0,0 +1,277 @@
+# coding=utf-8
+# Copyright 2024 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 List, Optional, Tuple, Union
+
+import numpy as np
+
+from transformers.processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+
+from ...image_processing_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...utils import is_vision_available, logging
+
+
+if is_vision_available():
+ from ...image_utils import load_images
+
+logger = logging.get_logger(__name__)
+
+
+class GotOcr2TextKwargs(TextKwargs, total=False):
+ format: Optional[bool]
+
+
+class GotOcr2ImagesKwargs(ImagesKwargs, total=False):
+ box: Optional[Union[List, Tuple[float, float], Tuple[float, float, float, float]]]
+ color: Optional[str]
+ num_image_tokens: Optional[int]
+ multi_page: Optional[bool]
+ crop_to_patches: Optional[bool]
+ min_patches: Optional[int]
+ max_patches: Optional[int]
+
+
+class GotOcr2ProcessorKwargs(ProcessingKwargs, total=False):
+ text_kwargs: GotOcr2TextKwargs
+ images_kwargs: GotOcr2ImagesKwargs
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ "format": False,
+ },
+ "images_kwargs": {
+ "num_image_tokens": 256,
+ "multi_page": False,
+ "crop_to_patches": False,
+ "min_patches": 1,
+ "max_patches": 12,
+ },
+ }
+
+
+def preprocess_box_annotation(box: Union[List, Tuple], image_size: Tuple[int, int]) -> List:
+ """
+ Convert box annotation to the format [x1, y1, x2, y2] in the range [0, 1000].
+ """
+ width, height = image_size
+ if len(box) == 4:
+ box[0] = int(box[0] / width * 1000)
+ box[1] = int(box[1] / height * 1000)
+ box[2] = int(box[2] / width * 1000)
+ box[3] = int(box[3] / height * 1000)
+ else:
+ raise ValueError("Box must be a list or tuple of lists in the form [x1, y1, x2, y2].")
+
+ return list(box)
+
+
+class GotOcr2Processor(ProcessorMixin):
+ r"""
+ Constructs a GotOcr2 processor which wraps a [`GotOcr2ImageProcessor`] and
+ [`PretrainedTokenizerFast`] tokenizer into a single processor that inherits both the image processor and
+ tokenizer functionalities. See the [`~GotOcr2Processor.__call__`] and [`~GotOcr2Processor.decode`] for more information.
+ Args:
+ image_processor ([`GotOcr2ImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`PreTrainedTokenizer`, `PreTrainedTokenizerFast`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ attributes = ["image_processor", "tokenizer"]
+ valid_kwargs = ["chat_template"]
+ image_processor_class = "AutoImageProcessor"
+ tokenizer_class = "PreTrainedTokenizerFast"
+
+ def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ self.message_start_token = "<|im_start|>"
+ self.message_end_token = "<|im_end|>"
+ self.img_start_token = "![]()
"
+ self.img_end_token = ""
+ self.img_pad_token = ""
+ self.system_query = "system\nYou should follow the instructions carefully and explain your answers in detail."
+
+ def _make_list_of_inputs(self, images, text, box, color, multi_page):
+ if not isinstance(images, (list, tuple)):
+ images = [images]
+ if multi_page:
+ logger.warning("Multi-page inference is enabled but only one image is passed.")
+ images = [images]
+ elif isinstance(images[0], (list, tuple)) and not multi_page:
+ raise ValueError("Nested images are only supported with `multi_page` set to `True`.")
+ elif not isinstance(images[0], (list, tuple)) and multi_page:
+ images = [images]
+
+ if isinstance(text, str):
+ text = [text]
+
+ if not isinstance(box[0], (list, tuple)):
+ # Use the same box for all images
+ box = [box for _ in range(len(images))]
+ if not isinstance(color, (list, tuple)):
+ color = [color for _ in range(len(images))]
+
+ return images, text, box, color
+
+ def __call__(
+ self,
+ images: Optional[ImageInput] = None,
+ text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None,
+ audio=None,
+ videos=None,
+ **kwargs: Unpack[GotOcr2ProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+ and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] to encode the text if `text`
+ is not `None`, otherwise encode default OCR queries which depends on the `format`, `box`, `color`, `multi_page` and
+ `crop_to_patches` arguments. To prepare the vision inputs, this method forwards the `images` and `kwrags` arguments to
+ GotOcr2ImageProcessor's [`~GotOcr2ImageProcessor.__call__`] if `images` is not `None`.
+
+ Args:
+ 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.
+ 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).
+ format (`bool`, *optional*):
+ If set, will add the format token to the query, and the model will return the OCR result with formatting.
+ box (`List[float]`, `List[Tuple[float, float]]`, `List[Tuple[float, float, float, float]]`, *optional*):
+ The box annotation to be added to the query. If a list of floats or a tuple of floats is provided, it
+ will be interpreted as [x1, y1, x2, y2]. If a list of tuples is provided, each tuple should be in the
+ form (x1, y1, x2, y2).
+ color (`str`, *optional*):
+ The color annotation to be added to the query. The model will return the OCR result within the box with
+ the specified color.
+ multi_page (`bool`, *optional*):
+ If set, will enable multi-page inference. The model will return the OCR result across multiple pages.
+ crop_to_patches (`bool`, *optional*):
+ If set, will crop the image to patches. The model will return the OCR result upon the patch reference.
+ min_patches (`int`, *optional*):
+ The minimum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to
+ `True`.
+ max_patches (`int`, *optional*):
+ The maximum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to
+ `True`.
+
+ 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:
+ [`BatchFeature`]: A [`BatchFeature`] 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`.
+ """
+
+ output_kwargs = self._merge_kwargs(
+ GotOcr2ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ format_output = output_kwargs["text_kwargs"].pop("format")
+ num_image_tokens = output_kwargs["images_kwargs"].pop("num_image_tokens")
+ box = output_kwargs["images_kwargs"].pop("box", [None])
+ color = output_kwargs["images_kwargs"].pop("color", None)
+ multi_page = output_kwargs["images_kwargs"].pop("multi_page")
+
+ crop_to_patches = output_kwargs["images_kwargs"].get("crop_to_patches")
+ images, text, box, color = self._make_list_of_inputs(images, text, box, color, multi_page)
+ if multi_page:
+ # save the number of pages per batch
+ num_pages_per_batch = [len(image_group) for image_group in images]
+ # flatten the list of images
+ images = [image for image_group in images for image in image_group]
+ else:
+ num_pages_per_batch = [1 for _ in range(len(images))]
+ # Load images as we need to know the image size
+ images = load_images(images)
+ image_sizes = [image.size for image in images]
+ image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ num_patches_array = image_inputs.pop("num_patches")
+ if text is None:
+ text = []
+ patch_indices = np.cumsum(num_pages_per_batch)
+ for index, (num_pages, box_single, color_single) in enumerate(zip(num_pages_per_batch, box, color)):
+ current_patch_index = patch_indices[index - 1] if index > 0 else 0
+ num_patches = sum(num_patches_array[current_patch_index : current_patch_index + num_pages])
+ if box_single[0] is not None:
+ box_single = preprocess_box_annotation(box_single, image_sizes[index])
+ query = (
+ f"{f'[{color_single}] ' if color_single is not None else ''}"
+ f"{str(box_single) if box_single[0] is not None else ''} "
+ "OCR"
+ f"{' with format' if format_output else ''}"
+ f"{' across multi pages' if multi_page else ''}"
+ f"{' upon the patch reference' if crop_to_patches else ''}"
+ ": "
+ )
+ prompt = (
+ self.message_start_token
+ + self.system_query
+ + self.message_end_token
+ + self.message_start_token
+ + "user\n"
+ + self.img_start_token
+ + self.img_pad_token * num_image_tokens * num_patches
+ + self.img_end_token
+ + "\n"
+ + query
+ + self.message_end_token
+ + self.message_start_token
+ + "assistant\n"
+ )
+ text.append(prompt)
+
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+ return BatchFeature(data={**text_inputs, **image_inputs})
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to PreTrainedTokenizerFast'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 PreTrainedTokenizerFast'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(tokenizer_input_names) + list(image_processor_input_names)
+
+
+__all__ = ["GotOcr2Processor"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__init__.py b/lib/python3.12/site-packages/transformers/models/gpt2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f01899e668e3a86548db3f59c7f42d70746385ab
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/__init__.py
@@ -0,0 +1,32 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_gpt2 import *
+ from .modeling_flax_gpt2 import *
+ from .modeling_gpt2 import *
+ from .modeling_tf_gpt2 import *
+ from .tokenization_gpt2 import *
+ from .tokenization_gpt2_fast import *
+ from .tokenization_gpt2_tf import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..146d8f7d441f7bf0f848191672d8421ea30cd6e8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2307ea3e864e5bf411c067a41b8b8edb4ba663c6
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..88a5dc9d28c571970f0965481ed58eed96b76963
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5791c0748918d358a09d80556c06cf85fd895b7e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7daa9822f9f294edbc8256960d68da7e421a0873
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..144dc3ce680f2d14e7e53651aa6b5a9095fa08a3
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9fb69d2c6d3cbda3106cbbd6a428685b3b197dfb
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..deda0d1171f5e834ebd081d2cafe28fdf90c6f16
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/configuration_gpt2.py b/lib/python3.12/site-packages/transformers/models/gpt2/configuration_gpt2.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3ebea02496c5057f72b5f1fdb90784f79d407fa
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/configuration_gpt2.py
@@ -0,0 +1,273 @@
+# 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__)
+
+
+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
+
+
+__all__ = ["GPT2Config", "GPT2OnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c52e32822b6342d3d4880ecf0cceaff949810ab
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py
@@ -0,0 +1,782 @@
+# 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 slightly 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,
+)
+
+
+__all__ = ["FlaxGPT2LMHeadModel", "FlaxGPT2Model", "FlaxGPT2PreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/modeling_gpt2.py b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_gpt2.py
new file mode 100644
index 0000000000000000000000000000000000000000..75c148f233acdce7e8363a550d6c06153f1fb3fa
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_gpt2.py
@@ -0,0 +1,1672 @@
+# 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 Callable, 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 ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa, _prepare_4d_causal_attention_mask_for_sdpa
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, 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,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.model_parallel_utils import assert_device_map, get_device_map
+from .configuration_gpt2 import GPT2Config
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "openai-community/gpt2"
+_CONFIG_FOR_DOC = "GPT2Config"
+
+
+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
+
+
+def eager_attention_forward(module, query, key, value, attention_mask, head_mask=None, **kwargs):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
+
+ if module.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 module.scale_attn_by_inverse_layer_idx:
+ attn_weights = attn_weights / float(module.layer_idx + 1)
+
+ if not module.is_cross_attention:
+ # if only "normal" attention layer implements causal mask
+ query_length, key_length = query.size(-2), key.size(-2)
+ causal_mask = module.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 = module.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)
+ attn_output = attn_output.transpose(1, 2)
+
+ return attn_output, attn_weights
+
+
+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 _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 torch.amp.autocast(query.device.type, 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, device=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)
+ attn_output = attn_output.transpose(1, 2)
+
+ return attn_output, attn_weights
+
+ 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,
+ **kwargs,
+ ) -> 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_states = self.q_attn(hidden_states)
+ key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
+ attention_mask = encoder_attention_mask
+ else:
+ query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
+
+ shape_q = (*query_states.shape[:-1], -1, self.head_dim)
+ shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
+
+ query_states = query_states.view(shape_q).transpose(1, 2)
+ key_states = key_states.view(shape_kv).transpose(1, 2)
+ value_states = value_states.view(shape_kv).transpose(1, 2)
+
+ if layer_past is not None:
+ past_key, past_value = layer_past
+ key_states = torch.cat((past_key, key_states), dim=-2)
+ value_states = torch.cat((past_value, value_states), dim=-2)
+
+ if use_cache is True:
+ present = (key_states, value_states)
+ else:
+ present = None
+
+ is_cross_attention = encoder_hidden_states is not None
+ is_causal = attention_mask is None and query_states.shape[-2] > 1 and not is_cross_attention
+
+ using_eager = self.config._attn_implementation == "eager"
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ if self.config._attn_implementation == "sdpa" and (output_attentions or head_mask is not None):
+ using_eager = True
+ logger.warning_once(
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ else:
+ # Attention functions are consistent with previous equivalent attention classes, however they do not support some options
+ # (e.g. layer scaling, head mask) that eager supports. These implementations are thus equivalent to previous code, but
+ # not necessarily to eager (if mentionned options are provided).
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ if using_eager and self.reorder_and_upcast_attn:
+ attn_output, attn_weights = self._upcast_and_reordered_attn(
+ query_states, key_states, value_states, attention_mask, head_mask
+ )
+ else:
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ head_mask=head_mask,
+ dropout=self.attn_dropout.p if self.training else 0.0,
+ is_causal=is_causal,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
+ 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 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
+
+
+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
+
+ self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ self.attn = GPT2Attention(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 = GPT2Attention(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
+ _supports_sdpa = 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: Optional[torch.FloatTensor] = None
+ mc_logits: Optional[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*):
+ 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):
+ _supports_param_buffer_assignment = False
+
+ 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)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.wte(input_ids)
+ position_embeds = self.wpe(position_ids)
+ hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)
+
+ # Attention mask.
+ _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
+ attention_mask = attention_mask.view(batch_size, -1) if attention_mask is not None else None
+ if self._attn_implementation == "flash_attention_2":
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ elif _use_sdpa:
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+ attention_mask=attention_mask,
+ input_shape=(batch_size, input_shape[-1]),
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_length,
+ )
+ else:
+ 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 = 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 _use_sdpa:
+ encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+ mask=encoder_attention_mask, dtype=inputs_embeds.dtype, tgt_len=input_shape[-1]
+ )
+ elif not 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 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 in range(len(self.h)):
+ block, layer_past = self.h[i], past_key_values[i]
+ # 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, GenerationMixin):
+ _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
+
+ @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,
+ **kwargs,
+ ) -> 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:
+ # Flatten the tokens
+ loss = self.loss_function(
+ lm_logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ 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, GenerationMixin):
+ _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
+
+ @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]
+
+ 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:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ 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), last_non_pad_token]
+
+ 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,
+ )
+
+
+__all__ = [
+ "GPT2DoubleHeadsModel",
+ "GPT2ForQuestionAnswering",
+ "GPT2ForSequenceClassification",
+ "GPT2ForTokenClassification",
+ "GPT2LMHeadModel",
+ "GPT2Model",
+ "GPT2PreTrainedModel",
+ "load_tf_weights_in_gpt2",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/modeling_tf_gpt2.py b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_tf_gpt2.py
new file mode 100644
index 0000000000000000000000000000000000000000..5812e42af75c5637091bc63c604fa76b44dbaee1
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/modeling_tf_gpt2.py
@@ -0,0 +1,1239 @@
+# 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"
+
+
+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: Optional[tf.Tensor] = None
+ mc_logits: Optional[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)
+ batch_size = logits_shape[0]
+
+ if self.config.pad_token_id is None:
+ last_non_pad_token = tf.fill((batch_size,), value=logits_shape[1] - 1)
+ else:
+ if input_ids is not None:
+ token_indices = tf.range(shape_list(input_ids)[-1])
+ non_pad_mask = tf.cast(input_ids != self.config.pad_token_id, token_indices.dtype)
+ last_non_pad_token = tf.reduce_max(token_indices * non_pad_mask, axis=-1)
+ else:
+ last_non_pad_token = tf.fill((batch_size,), value=logits_shape[1] - 1)
+ logger.warning_once(
+ 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
+
+ pooled_logits = tf.gather(logits, last_non_pad_token, batch_dims=1, axis=1)
+
+ 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.")
+
+ loss = self.hf_compute_loss(tf.reshape(labels, [-1]), tf.reshape(pooled_logits, [-1, self.num_labels]))
+
+ 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)
+
+
+__all__ = [
+ "TFGPT2DoubleHeadsModel",
+ "TFGPT2ForSequenceClassification",
+ "TFGPT2LMHeadModel",
+ "TFGPT2MainLayer",
+ "TFGPT2Model",
+ "TFGPT2PreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2.py b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2.py
new file mode 100644
index 0000000000000000000000000000000000000000..709bcec5b611231ea8bfc3d72322b0ebeb1a9a1e
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2.py
@@ -0,0 +1,334 @@
+# 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)
+
+
+__all__ = ["GPT2Tokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..81e67a818deebeeeed8f9020510c9fbb6e0c1d5e
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py
@@ -0,0 +1,133 @@
+# 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."""
+
+from typing import Optional, Tuple
+
+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=vocab_file,
+ merges_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)
+
+ 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)
+
+
+__all__ = ["GPT2TokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_tf.py b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_tf.py
new file mode 100644
index 0000000000000000000000000000000000000000..34e6ca2d2555c5a357c3bcdef1e705b83e626e0a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt2/tokenization_gpt2_tf.py
@@ -0,0 +1,113 @@
+import os
+from typing import Dict, List, Optional, 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: Optional[int] = None,
+ pad_token_id: Optional[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: Optional[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}
+
+
+__all__ = ["TFGPT2Tokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__init__.py b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..92e985d92734550a5b0635941294669386d35749
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_gpt_bigcode import *
+ from .modeling_gpt_bigcode import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b1b00fef235c7e25c9fc90eda2425e91acc05b5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eeec728f0d8e016f467412a71824ddceb95bab8b
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f81e237a674e3ef7864c79c958fa8c4879a342c8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py
new file mode 100644
index 0000000000000000000000000000000000000000..46a3dfea44101fb07666a82bad4a127137486f99
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py
@@ -0,0 +1,144 @@
+# 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__)
+
+
+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)
+
+
+__all__ = ["GPTBigCodeConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
new file mode 100644
index 0000000000000000000000000000000000000000..d9a84d6f463d979aa9b3c45e36ffe0ba6d73966d
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
@@ -0,0 +1,1428 @@
+# 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.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+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,
+ logging,
+)
+from .configuration_gpt_bigcode import GPTBigCodeConfig
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "bigcode/gpt_bigcode-santacoder"
+_CONFIG_FOR_DOC = "GPTBigCodeConfig"
+
+
+# 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
+
+
+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.
+ """
+
+ 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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ 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 = _flash_attention_forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ query_length,
+ dropout=attn_dropout,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ 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)
+
+
+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()
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ # 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 = True if self.is_causal and attention_mask is None and query_length > 1 else False
+
+ 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,
+ is_causal=is_causal,
+ 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,
+ **kwargs,
+ ) -> 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 in ["cuda", "xpu"]:
+ # 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.to(inputs_embeds.device)
+
+ 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, GenerationMixin):
+ _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):
+ # Overwritten -- `past_key_values` with uncommon shape
+
+ 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
+
+ def _get_initial_cache_position(self, input_ids, model_kwargs):
+ """
+ Calculates `cache_position` for the pre-fill stage based on `input_ids` and optionally past length.
+ Since gpt bigcode is special, the method is overridden here, other models use it from `generation.utils.py`.
+ """
+ past_length = 0
+ if "past_key_values" in model_kwargs:
+ if self.config.multi_query:
+ past_length = model_kwargs["past_key_values"][0].shape[1]
+ else:
+ past_length = model_kwargs["past_key_values"][0].shape[2]
+ if "inputs_embeds" in model_kwargs:
+ cur_len = model_kwargs["inputs_embeds"].shape[1]
+ else:
+ cur_len = input_ids.shape[-1]
+ model_kwargs["cache_position"] = torch.arange(past_length, cur_len, device=input_ids.device)
+ return model_kwargs
+
+ @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,
+ **kwargs,
+ ) -> 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:
+ loss = self.loss_function(
+ lm_logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ 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]
+
+ 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:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ 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), last_non_pad_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(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,
+ )
+
+
+__all__ = [
+ "GPTBigCodeForSequenceClassification",
+ "GPTBigCodeForTokenClassification",
+ "GPTBigCodeForCausalLM",
+ "GPTBigCodeModel",
+ "GPTBigCodePreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/granite/__init__.py b/lib/python3.12/site-packages/transformers/models/granite/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a98daa072d583618634afb0393e17860c2f1149
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/granite/__init__.py
@@ -0,0 +1,57 @@
+# Copyright 2024 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_torch_available,
+)
+
+
+_import_structure = {
+ "configuration_granite": ["GraniteConfig"],
+}
+
+try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ pass
+else:
+ _import_structure["modeling_granite"] = [
+ "GraniteForCausalLM",
+ "GraniteModel",
+ "GranitePreTrainedModel",
+ ]
+
+if TYPE_CHECKING:
+ from .configuration_granite import GraniteConfig
+
+ try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ pass
+ else:
+ from .modeling_granite import (
+ GraniteForCausalLM,
+ GraniteModel,
+ GranitePreTrainedModel,
+ )
+
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/granite/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..411008314c134bc836a55fdad581bc4190fef91c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/granite/__pycache__/configuration_granite.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/configuration_granite.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3c67e069ec9c3a1f83ca4ac3cd59710443d88875
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/configuration_granite.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modeling_granite.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modeling_granite.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5cb37e22c8c4edd084533c303f45021159742627
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modeling_granite.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modular_granite.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modular_granite.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..466a56c6e805befc42afea70b9865752bab632d4
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/granite/__pycache__/modular_granite.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/granite/configuration_granite.py b/lib/python3.12/site-packages/transformers/models/granite/configuration_granite.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc651a94e1bdd8901cb89365984538591669c469
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/granite/configuration_granite.py
@@ -0,0 +1,194 @@
+# coding=utf-8
+# Copyright 2024 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.
+"""Granite model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...modeling_rope_utils import rope_config_validation
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class GraniteConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`GraniteModel`]. It is used to instantiate an Granite
+ 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 Granite-3B.
+
+ 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 Granite model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`GraniteModel`]
+ 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 decoder.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ 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 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.
+ 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*):
+ Padding token id.
+ bos_token_id (`int`, *optional*, defaults to 1):
+ Beginning of stream token id.
+ eos_token_id (`int`, *optional*, defaults to 2):
+ End of stream token id.
+ 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.
+ attention_bias (`bool`, *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.
+ mlp_bias (`bool`, *optional*, defaults to `False`):
+ Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
+ embedding_multiplier (`float`, *optional*, defaults to 1.0): embedding multiplier
+ logits_scaling (`float`, *optional*, defaults to 1.0): divisor for output logits
+ residual_multiplier (`float`, *optional*, defaults to 1.0): residual multiplier
+ attention_multiplier (`float`, *optional*, defaults to 1.0): attention multiplier
+
+ ```python
+ >>> from transformers import GraniteModel, GraniteConfig
+
+ >>> # Initializing a Granite granite-3b style configuration
+ >>> configuration = GraniteConfig()
+
+ >>> # Initializing a model from the granite-7b style configuration
+ >>> model = GraniteModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "granite"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ # Default tensor parallel plan for base model `GraniteModel`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=32000,
+ hidden_size=4096,
+ intermediate_size=11008,
+ num_hidden_layers=32,
+ num_attention_heads=32,
+ num_key_value_heads=None,
+ hidden_act="silu",
+ max_position_embeddings=2048,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ pad_token_id=None,
+ bos_token_id=1,
+ eos_token_id=2,
+ tie_word_embeddings=False,
+ rope_theta=10000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ attention_dropout=0.0,
+ mlp_bias=False,
+ embedding_multiplier=1.0,
+ logits_scaling=1.0,
+ residual_multiplier=1.0,
+ attention_multiplier=1.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
+
+ # 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.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ self.mlp_bias = mlp_bias
+
+ self.embedding_multiplier = embedding_multiplier
+ self.logits_scaling = logits_scaling
+ self.residual_multiplier = residual_multiplier
+ self.attention_multiplier = attention_multiplier
+
+ 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,
+ )
+
+ rope_config_validation(self)
diff --git a/lib/python3.12/site-packages/transformers/models/granite/modeling_granite.py b/lib/python3.12/site-packages/transformers/models/granite/modeling_granite.py
new file mode 100644
index 0000000000000000000000000000000000000000..e65f43229146af41903fbdb4653d6d9ce1817418
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/granite/modeling_granite.py
@@ -0,0 +1,860 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/granite/modular_granite.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_granite.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 IBM 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 functools import partial
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ LossKwargs,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ can_return_tuple,
+ is_torch_flex_attn_available,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.deprecation import deprecate_kwarg
+from .configuration_granite import GraniteConfig
+
+
+if is_torch_flex_attn_available():
+ from torch.nn.attention.flex_attention import BlockMask
+
+ from ...integrations.flex_attention import make_flex_block_causal_mask
+
+
+logger = logging.get_logger(__name__)
+_CONFIG_FOR_DOC = "GraniteConfig"
+
+
+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, 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
+
+
+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 eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class GraniteAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: GraniteConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = config.attention_multiplier
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_value: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ 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; 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)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+ logger.warning_once(
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ else:
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GraniteRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ GraniteRMSNorm 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class GraniteMLP(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=config.mlp_bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class GraniteDecoderLayer(nn.Module):
+ def __init__(self, config: GraniteConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = GraniteAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = GraniteMLP(config)
+ self.input_layernorm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.residual_multiplier = config.residual_multiplier
+
+ 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: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **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
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence
+ position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights = 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,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states * self.residual_multiplier
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states * self.residual_multiplier # main diff with Llama
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights,)
+
+ return outputs
+
+
+class GraniteRotaryEmbedding(nn.Module):
+ def __init__(self, config: GraniteConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+GRANITE_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 ([`GraniteConfig`]):
+ 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 Granite Model outputting raw hidden-states without any specific head on top.",
+ GRANITE_START_DOCSTRING,
+)
+class GranitePreTrainedModel(PreTrainedModel):
+ config_class = GraniteConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GraniteDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _supports_cache_class = True
+ _supports_quantized_cache = True
+ _supports_static_cache = True
+ _supports_attention_backend = 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_()
+
+
+GRANITE_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`, *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`.
+
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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 Granite Model outputting raw hidden-states without any specific head on top.",
+ GRANITE_START_DOCSTRING,
+)
+class GraniteModel(GranitePreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GraniteDecoderLayer`]
+
+ Args:
+ config: GraniteConfig
+ """
+
+ def __init__(self, config: GraniteConfig):
+ 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(
+ [GraniteDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = GraniteRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = GraniteRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.embedding_multiplier = config.embedding_multiplier
+
+ # 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
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(GRANITE_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.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> 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
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ 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)
+
+ inputs_embeds = inputs_embeds * self.embedding_multiplier # main diff with Llama
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache()
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ 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_key_values, output_attentions
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ partial(decoder_layer.__call__, **flash_attn_kwargs),
+ hidden_states,
+ causal_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ cache_position,
+ position_embeddings,
+ )
+ 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,
+ position_embeddings=position_embeddings,
+ **flash_attn_kwargs,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ 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,)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ 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_key_values: Cache,
+ output_attentions: bool = False,
+ ):
+ if self.config._attn_implementation == "flash_attention_2":
+ if attention_mask is not None and (attention_mask == 0.0).any():
+ return attention_mask
+ return None
+ if self.config._attn_implementation == "flex_attention":
+ if isinstance(attention_mask, torch.Tensor):
+ attention_mask = make_flex_block_causal_mask(attention_mask)
+ if isinstance(attention_mask, BlockMask):
+ return attention_mask
+
+ # 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. This feature is not compatible with static cache, as SDPA will fail
+ # to infer the attention mask.
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ using_static_cache = isinstance(past_key_values, StaticCache)
+
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
+ attention_mask,
+ inputs_embeds=input_tensor,
+ past_key_values_length=past_seen_tokens,
+ is_training=self.training,
+ ):
+ return None
+
+ dtype, device = input_tensor.dtype, input_tensor.device
+ sequence_length = input_tensor.shape[1]
+ if using_static_cache:
+ target_length = past_key_values.get_max_cache_shape()
+ else:
+ target_length = (
+ attention_mask.shape[-1]
+ if isinstance(attention_mask, torch.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask,
+ sequence_length=sequence_length,
+ target_length=target_length,
+ dtype=dtype,
+ device=device,
+ cache_position=cache_position,
+ batch_size=input_tensor.shape[0],
+ )
+
+ if (
+ self.config._attn_implementation == "sdpa"
+ and attention_mask is not None
+ and attention_mask.device.type in ["cuda", "xpu"]
+ and not output_attentions
+ ):
+ # 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
+ min_dtype = torch.finfo(dtype).min
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+ return causal_mask
+
+ @staticmethod
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ **kwargs,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+ `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache,
+ to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ min_dtype = torch.finfo(dtype).min
+ 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(batch_size, 1, -1, -1)
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+ causal_mask.device
+ )
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+ _tp_plan = {"lm_head": "colwise_rep"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = GraniteModel(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
+
+ @can_return_tuple
+ @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+ @add_start_docstrings_to_model_forward(GRANITE_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, 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[Union[Cache, 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,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[KwargsForCausalLM],
+ ) -> 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, GraniteForCausalLM
+
+ >>> model = GraniteForCausalLM.from_pretrained("meta-granite/Granite-2-7b-hf")
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-granite/Granite-2-7b-hf")
+
+ >>> 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
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = 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,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ logits = logits / self.config.logits_scaling # main diff with Llama
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
diff --git a/lib/python3.12/site-packages/transformers/models/granite/modular_granite.py b/lib/python3.12/site-packages/transformers/models/granite/modular_granite.py
new file mode 100644
index 0000000000000000000000000000000000000000..25929dbb33765d9e86d99f4bd1b4036f44814ee0
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/granite/modular_granite.py
@@ -0,0 +1,283 @@
+# coding=utf-8
+# Copyright 2024 IBM 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 functools import partial
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...cache_utils import Cache, DynamicCache
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...processing_utils import Unpack
+from ...utils import LossKwargs, logging
+from ..llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel
+from .configuration_granite import GraniteConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class GraniteAttention(LlamaAttention):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: GraniteConfig, layer_idx: Optional[int] = None):
+ super().__init__(config, layer_idx)
+ self.scaling = config.attention_multiplier
+
+
+class GraniteDecoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: GraniteConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.residual_multiplier = config.residual_multiplier
+ self.self_attn = GraniteAttention(config=config, layer_idx=layer_idx)
+
+ 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: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **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
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence
+ position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights = 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,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states * self.residual_multiplier
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states * self.residual_multiplier # main diff with Llama
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights,)
+
+ return outputs
+
+
+class GraniteModel(LlamaModel):
+ def __init__(self, config: GraniteConfig):
+ super().__init__(config)
+ self.embedding_multiplier = config.embedding_multiplier
+ self.layers = nn.ModuleList(
+ [GraniteDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+ 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.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> 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
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ 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)
+
+ inputs_embeds = inputs_embeds * self.embedding_multiplier # main diff with Llama
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache()
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ 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_key_values, output_attentions
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ partial(decoder_layer.__call__, **flash_attn_kwargs),
+ hidden_states,
+ causal_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ cache_position,
+ position_embeddings,
+ )
+ 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,
+ position_embeddings=position_embeddings,
+ **flash_attn_kwargs,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ 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,)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+class GraniteForCausalLM(LlamaForCausalLM):
+ 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[Union[Cache, 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,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[KwargsForCausalLM],
+ ) -> CausalLMOutputWithPast:
+ 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
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = 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,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ logits = logits / self.config.logits_scaling # main diff with Llama
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
diff --git a/lib/python3.12/site-packages/transformers/models/helium/__init__.py b/lib/python3.12/site-packages/transformers/models/helium/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..73d0966e5c1c489944c4b539c0cc06384c985c87
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/helium/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_helium import *
+ from .modeling_helium import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/helium/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c9d28a4a385b2d1f9246887e9b716ec3f4f7dde8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/helium/__pycache__/configuration_helium.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/configuration_helium.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9eb57f90966ce363f124c5c7c1cc73ae7a85f42a
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/configuration_helium.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modeling_helium.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modeling_helium.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..482cb064374107dd8b2f927c565471f690221959
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modeling_helium.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modular_helium.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modular_helium.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..435e1f7e8e46dfbf76d95d0bee3794cc40e11eaf
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/helium/__pycache__/modular_helium.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/helium/configuration_helium.py b/lib/python3.12/site-packages/transformers/models/helium/configuration_helium.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b27c6e54b69492ea805540ae2d50aa61ceb2617
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/helium/configuration_helium.py
@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2024 The Kyutai and HuggingFace Inc. 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 ...configuration_utils import PretrainedConfig
+
+
+class HeliumConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`HeliumModel`]. It is used to instantiate an Helium
+ 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 Helium 2b model.
+ e.g. [kyutai/helium-2b](https://huggingface.co/kyutai/helium-2b)
+ 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 48000):
+ Vocabulary size of the Helium model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`HeliumModel`]
+ hidden_size (`int`, *optional*, defaults to 2560):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 7040):
+ Dimension of the MLP representations.
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of hidden layers in the Transformer decoder.
+ num_attention_heads (`int`, *optional*, defaults to 20):
+ Number of attention heads for each attention layer in the Transformer decoder.
+ num_key_value_heads (`int`, *optional*, defaults to 20):
+ 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 128):
+ The attention head dimension.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The legacy activation function. It is overwritten by the `hidden_activation`.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
+ 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-08):
+ 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 100000.0):
+ The base period of the RoPE embeddings.
+ pad_token_id (`int`, *optional*, defaults to 3):
+ Padding token id.
+ eos_token_id (`int` | `list`, *optional*, defaults to 2):
+ End of stream token id.
+ bos_token_id (`int`, *optional*, defaults to 1):
+ Beginning of stream token id.
+ attention_bias (`bool`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ mlp_bias (`bool`, *optional*, defaults to `False`):
+ Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
+ ```python
+ >>> from transformers import HeliumModel, HeliumConfig
+ >>> # Initializing a Helium 2b style configuration
+ >>> configuration = HeliumConfig()
+ >>> # Initializing a model from the Helium 2b style configuration
+ >>> model = HeliumModel(configuration)
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "helium"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=48000,
+ hidden_size=2560,
+ intermediate_size=7040,
+ num_hidden_layers=24,
+ num_attention_heads=20,
+ num_key_value_heads=20,
+ head_dim=128,
+ hidden_act="silu",
+ attention_dropout=0.0,
+ max_position_embeddings=4096,
+ initializer_range=0.02,
+ rms_norm_eps=1e-8,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=100000.0,
+ pad_token_id=3,
+ eos_token_id=2,
+ bos_token_id=1,
+ attention_bias=False,
+ mlp_bias=False,
+ **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.num_key_value_heads = num_key_value_heads
+ self.head_dim = head_dim
+ self.hidden_act = hidden_act
+ 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
+ self.mlp_bias = mlp_bias
+
+ 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,
+ )
+
+
+__all__ = ["HeliumConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/helium/modeling_helium.py b/lib/python3.12/site-packages/transformers/models/helium/modeling_helium.py
new file mode 100644
index 0000000000000000000000000000000000000000..37350ae4626b3febe8e846017d8bc7e37999016f
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/helium/modeling_helium.py
@@ -0,0 +1,1039 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/helium/modular_helium.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_helium.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2024 The Kyutai and HuggingFace Inc. 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.
+import math
+from functools import partial
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import (
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ LossKwargs,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ can_return_tuple,
+ is_torch_flex_attn_available,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.deprecation import deprecate_kwarg
+from .configuration_helium import HeliumConfig
+
+
+if is_torch_flex_attn_available():
+ from torch.nn.attention.flex_attention import BlockMask
+
+ from ...integrations.flex_attention import make_flex_block_causal_mask
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "google/helium-7b"
+_CONFIG_FOR_DOC = "HeliumConfig"
+
+
+class HeliumRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ 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.to(torch.float32) * hidden_states).to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class HeliumRotaryEmbedding(nn.Module):
+ def __init__(self, config: HeliumConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class HeliumMLP(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=config.mlp_bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+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 eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+ return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+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)
+
+ # Interleave them instead of usual shape
+ cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
+ sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+
+ return q_embed, k_embed
+
+
+class HeliumAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: HeliumConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = 1 / math.sqrt(self.head_dim)
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_value: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ 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; 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)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+ logger.warning_once(
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ else:
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class HeliumDecoderLayer(nn.Module):
+ def __init__(self, config: HeliumConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = HeliumAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = HeliumMLP(config)
+ self.input_layernorm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = HeliumRMSNorm(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[Cache] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights = 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,
+ position_embeddings=position_embeddings,
+ **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,)
+
+ return outputs
+
+
+HELIUM_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 ([`HeliumConfig`]):
+ 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 Helium Model outputting raw hidden-states without any specific head on top.",
+ HELIUM_START_DOCSTRING,
+)
+class HeliumPreTrainedModel(PreTrainedModel):
+ config_class = HeliumConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["HeliumDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _supports_cache_class = True
+ _supports_quantized_cache = True
+ _supports_static_cache = True
+ _supports_attention_backend = 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_()
+
+
+HELIUM_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`, *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`.
+
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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 Helium Model outputting raw hidden-states without any specific head on top.",
+ HELIUM_START_DOCSTRING,
+)
+class HeliumModel(HeliumPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`HeliumDecoderLayer`]
+
+ Args:
+ config: HeliumConfig
+ """
+
+ def __init__(self, config: HeliumConfig):
+ 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(
+ [HeliumDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = HeliumRotaryEmbedding(config)
+ 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
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(HELIUM_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.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> 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
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ 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
+
+ # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
+ if not isinstance(past_key_values, (type(None), Cache)):
+ raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache()
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ 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_key_values, output_attentions
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ partial(decoder_layer.__call__, **flash_attn_kwargs),
+ hidden_states,
+ causal_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ cache_position,
+ position_embeddings,
+ )
+ 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,
+ position_embeddings=position_embeddings,
+ **flash_attn_kwargs,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ 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,)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ 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_key_values: Cache,
+ output_attentions: bool = False,
+ ):
+ if self.config._attn_implementation == "flash_attention_2":
+ if attention_mask is not None and (attention_mask == 0.0).any():
+ return attention_mask
+ return None
+ if self.config._attn_implementation == "flex_attention":
+ if isinstance(attention_mask, torch.Tensor):
+ attention_mask = make_flex_block_causal_mask(attention_mask)
+ if isinstance(attention_mask, BlockMask):
+ return attention_mask
+
+ # 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. This feature is not compatible with static cache, as SDPA will fail
+ # to infer the attention mask.
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ using_static_cache = isinstance(past_key_values, StaticCache)
+
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
+ attention_mask,
+ inputs_embeds=input_tensor,
+ past_key_values_length=past_seen_tokens,
+ is_training=self.training,
+ ):
+ return None
+
+ dtype, device = input_tensor.dtype, input_tensor.device
+ sequence_length = input_tensor.shape[1]
+ if using_static_cache:
+ target_length = past_key_values.get_max_cache_shape()
+ else:
+ target_length = (
+ attention_mask.shape[-1]
+ if isinstance(attention_mask, torch.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask,
+ sequence_length=sequence_length,
+ target_length=target_length,
+ dtype=dtype,
+ device=device,
+ cache_position=cache_position,
+ batch_size=input_tensor.shape[0],
+ )
+
+ if (
+ self.config._attn_implementation == "sdpa"
+ and attention_mask is not None
+ and attention_mask.device.type in ["cuda", "xpu"]
+ and not output_attentions
+ ):
+ # 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
+ min_dtype = torch.finfo(dtype).min
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+ return causal_mask
+
+ @staticmethod
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ **kwargs,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+ `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache,
+ to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ min_dtype = torch.finfo(dtype).min
+ 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(batch_size, 1, -1, -1)
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+ causal_mask.device
+ )
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+ _tp_plan = {"lm_head": "colwise_rep"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.model = HeliumModel(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
+
+ @can_return_tuple
+ @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+ @add_start_docstrings_to_model_forward(HELIUM_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, 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.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[KwargsForCausalLM],
+ ) -> 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, HeliumForCausalLM
+
+ >>> model = HeliumForCausalLM.from_pretrained("google/helium-7b")
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/helium-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
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = 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,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ The Helium Model transformer with a sequence classification head on top (linear layer).
+
+ [`HeliumForSequenceClassification`] 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).
+ """,
+ HELIUM_START_DOCSTRING,
+)
+class HeliumForSequenceClassification(HeliumPreTrainedModel):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = HeliumModel(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
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(HELIUM_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.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ ) -> 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).
+ """
+
+ transformer_outputs: BaseModelOutputWithPast = 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,
+ )
+ hidden_states = transformer_outputs.last_hidden_state
+ 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:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ 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), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+ 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 Helium Model transformer 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.
+ """,
+ HELIUM_START_DOCSTRING,
+)
+class HeliumForTokenClassification(HeliumPreTrainedModel):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = HeliumModel(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.score = nn.Linear(config.hidden_size, config.num_labels)
+
+ # 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
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(HELIUM_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,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = 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,
+ ) -> 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).
+ """
+
+ outputs: BaseModelOutputWithPast = 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,
+ )
+ sequence_output = outputs.last_hidden_state
+ sequence_output = self.dropout(sequence_output)
+ logits = self.score(sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.config)
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "HeliumPreTrainedModel",
+ "HeliumModel",
+ "HeliumForCausalLM",
+ "HeliumForSequenceClassification",
+ "HeliumForTokenClassification",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/helium/modular_helium.py b/lib/python3.12/site-packages/transformers/models/helium/modular_helium.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c9be5ec805afd5b04530ab4b3dda1500c81689f
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/helium/modular_helium.py
@@ -0,0 +1,171 @@
+# coding=utf-8
+# Copyright 2024 The Kyutai and HuggingFace Inc. 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.
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from ...utils import logging
+from ..gemma.modeling_gemma import (
+ GemmaForCausalLM,
+ GemmaForSequenceClassification,
+ GemmaForTokenClassification,
+)
+from ..granite.modeling_granite import (
+ GraniteAttention,
+)
+from ..llama.modeling_llama import (
+ LlamaDecoderLayer,
+ LlamaMLP,
+ LlamaModel,
+ LlamaPreTrainedModel,
+ LlamaRotaryEmbedding,
+)
+from .configuration_helium import HeliumConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class HeliumRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ 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.to(torch.float32) * hidden_states).to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class HeliumRotaryEmbedding(LlamaRotaryEmbedding):
+ pass
+
+
+class HeliumMLP(LlamaMLP):
+ pass
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+ return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+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)
+
+ # Interleave them instead of usual shape
+ cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
+ sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+
+ return q_embed, k_embed
+
+
+class HeliumAttention(GraniteAttention):
+ def __init__(self, config: HeliumConfig, layer_idx: Optional[int] = None):
+ super().__init__(config, layer_idx)
+ self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+ self.scaling = 1 / math.sqrt(self.head_dim)
+
+
+class HeliumDecoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: HeliumConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+
+ self.mlp = HeliumMLP(config)
+ self.input_layernorm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+class HeliumPreTrainedModel(LlamaPreTrainedModel):
+ pass
+
+
+class HeliumModel(HeliumPreTrainedModel, LlamaModel):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.layers = nn.ModuleList(
+ [HeliumDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = HeliumRotaryEmbedding(config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+
+class HeliumForCausalLM(GemmaForCausalLM):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.model = HeliumModel(config)
+ self.post_init()
+
+
+class HeliumForSequenceClassification(GemmaForSequenceClassification):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.model = HeliumModel(config)
+ self.post_init()
+
+
+class HeliumForTokenClassification(GemmaForTokenClassification):
+ def __init__(self, config: HeliumConfig):
+ super().__init__(config)
+ self.model = HeliumModel(config)
+ self.post_init()
+
+
+__all__ = [
+ "HeliumPreTrainedModel",
+ "HeliumModel",
+ "HeliumForCausalLM",
+ "HeliumForSequenceClassification",
+ "HeliumForTokenClassification",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/__init__.py b/lib/python3.12/site-packages/transformers/models/jetmoe/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7058590acc8859f934330ced9556d6ca66b50a51
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/jetmoe/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_jetmoe import *
+ from .modeling_jetmoe import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..00e55b06a76b78636799f3788aabf9c5a9470782
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/configuration_jetmoe.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/configuration_jetmoe.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..78827c2363b693a719356a4675059d1956e4bd75
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/configuration_jetmoe.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/modeling_jetmoe.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/modeling_jetmoe.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..178961ae05b08a37b4149c7e50647612b66404b7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/jetmoe/__pycache__/modeling_jetmoe.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/configuration_jetmoe.py b/lib/python3.12/site-packages/transformers/models/jetmoe/configuration_jetmoe.py
new file mode 100644
index 0000000000000000000000000000000000000000..4053c900bf2dc84eb86d8c03ac1cd94cc2609a07
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/jetmoe/configuration_jetmoe.py
@@ -0,0 +1,152 @@
+# coding=utf-8
+# Copyright 2024 JetMoe 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.
+"""JetMoe model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class JetMoeConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`JetMoeModel`]. It is used to instantiate a
+ JetMoe model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a configuration of the JetMoe-4B.
+
+ [jetmoe/jetmoe-8b](https://huggingface.co/jetmoe/jetmoe-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 32000):
+ Vocabulary size of the JetMoe model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`JetMoeModel`]
+ hidden_size (`int`, *optional*, defaults to 2048):
+ Dimension of the hidden representations.
+ num_hidden_layers (`int`, *optional*, defaults to 12):
+ Number of hidden layers in the Transformer encoder.
+ num_key_value_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads for each key and value in the Transformer encoder.
+ kv_channels (`int`, *optional*, defaults to 128):
+ Defines the number of channels for the key and value tensors.
+ intermediate_size (`int`, *optional*, defaults to 5632):
+ Dimension of the MLP representations.
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
+ The maximum sequence length that this model might ever be used with. JetMoe's attention allows sequence of
+ up to 4096 tokens.
+ activation_function (`string`, *optional*, defaults to `"silu"`):
+ Defines the activation function for MLP experts.
+ num_local_experts (`int`, *optional*, defaults to 8):
+ Defines the number of experts in the MoE and MoA.
+ num_experts_per_tok (`int, *optional*, defaults to 2):
+ The number of experts to route per-token and for MoE and MoA.
+ output_router_logits (`bool`, *optional*, defaults to `False`):
+ Whether or not the router logits should be returned by the model. Enabeling this will also
+ allow the model to output the auxiliary loss.
+ aux_loss_coef (`float`, *optional*, defaults to 0.01):
+ The coefficient for the auxiliary loss.
+ 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 1):
+ The id of the "beginning-of-sequence" token.
+ eos_token_id (`int`, *optional*, defaults to 2):
+ The id of the "end-of-sequence" token.
+ tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+ Whether the model's input and output word embeddings should be tied.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ initializer_range (`float`, *optional*, defaults to 0.01):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+
+ ```python
+ >>> from transformers import JetMoeModel, JetMoeConfig
+
+ >>> # Initializing a JetMoe 4B style configuration
+ >>> configuration = JetMoeConfig()
+
+ >>> # Initializing a model from the JetMoe 4B style configuration
+ >>> model = JetMoeModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "jetmoe"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=32000,
+ hidden_size=2048,
+ num_hidden_layers=12,
+ num_key_value_heads=16,
+ kv_channels=128,
+ intermediate_size=5632,
+ max_position_embeddings=4096,
+ activation_function="silu",
+ num_local_experts=8,
+ num_experts_per_tok=2,
+ output_router_logits=False,
+ aux_loss_coef=0.01,
+ use_cache=True,
+ bos_token_id=1,
+ eos_token_id=2,
+ tie_word_embeddings=True,
+ rope_theta=10000.0,
+ rms_norm_eps=1e-6,
+ initializer_range=0.01,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ if num_experts_per_tok > num_local_experts:
+ raise ValueError("`num_experts_per_tok` must be less than or equal to `num_local_experts`")
+ self.vocab_size = vocab_size
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_key_value_heads * num_experts_per_tok
+ self.num_key_value_heads = num_key_value_heads
+ self.kv_channels = kv_channels
+ self.intermediate_size = intermediate_size
+ self.max_position_embeddings = max_position_embeddings
+ self.activation_function = activation_function
+ self.num_local_experts = num_local_experts
+ self.num_experts_per_tok = num_experts_per_tok
+ self.output_router_logits = output_router_logits
+ self.aux_loss_coef = aux_loss_coef
+ self.use_cache = use_cache
+ self.initializer_range = initializer_range
+ self.attention_dropout = attention_dropout
+
+ self.bos_token_id = bos_token_id
+ self.eos_token_id = eos_token_id
+
+ self.rope_theta = rope_theta
+ self.rms_norm_eps = rms_norm_eps
+
+ super().__init__(
+ bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
+ )
+
+
+__all__ = ["JetMoeConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/jetmoe/modeling_jetmoe.py b/lib/python3.12/site-packages/transformers/models/jetmoe/modeling_jetmoe.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a4b6b36f1a7e4ea3081d4c5d5c640bdc70a540a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/jetmoe/modeling_jetmoe.py
@@ -0,0 +1,1465 @@
+# coding=utf-8
+# Copyright 2024 JetMoe 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 JetMoe model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import functional as F
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_outputs import (
+ MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+ SequenceClassifierOutputWithPast,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ can_return_tuple,
+ is_torch_flex_attn_available,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.deprecation import deprecate_kwarg
+from .configuration_jetmoe import JetMoeConfig
+
+
+if is_torch_flex_attn_available():
+ from torch.nn.attention.flex_attention import BlockMask
+
+ from ...integrations.flex_attention import make_flex_block_causal_mask
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "jetmoe"
+_CONFIG_FOR_DOC = "JetMoeConfig"
+
+
+# Copied from transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func
+def load_balancing_loss_func(
+ gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None],
+ num_experts: Optional[int] = None,
+ top_k=2,
+ attention_mask: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, int]:
+ 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:
+ 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:
+ Number of experts
+ top_k:
+ The number of experts to route per-token, can be also interpreted as the `top-k` routing
+ parameter.
+ attention_mask (`torch.Tensor`, *optional*):
+ 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 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
+
+
+class JetMoeParallelExperts(nn.Module):
+ def __init__(self, num_experts: int, input_size: int, output_size: int) -> None:
+ """
+ Initialize the JetMoeParallelExperts module.
+ The experts weights are stored in [num_experts, output_size, input_size] format. Such that it's comptible with
+ many MoE libraries, such as [Megablock](https://github.com/databricks/megablocks) and
+ [ScatterMoE](https://github.com/shawntan/scattermoe), as well as the
+ [MoE kernel](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py)
+ used in vllm.
+
+ Args:
+ num_experts (int):
+ Number of experts.
+ input_size (int):
+ Size of the input.
+ output_size (int):
+ Size of the output.
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size))
+ self.num_experts = num_experts
+ self.input_size = input_size
+ self.output_size = output_size
+
+ def forward(self, inputs, expert_size):
+ """
+ Forward pass of the JetMoeParallelExperts module.
+
+ Args:
+ inputs (Tensor):
+ Input tensor.
+ expert_size:
+ Expert size information.
+
+ Returns:
+ Tensor: Output tensor.
+ """
+ input_list = inputs.split(expert_size, dim=0)
+ output_list = []
+ for i in range(self.num_experts):
+ output_list.append(F.linear(input_list[i], self.weight[i]))
+ results = torch.cat(output_list, dim=0)
+ return results
+
+
+class JetMoeTopKGating(nn.Module):
+ def __init__(self, input_size: int, num_experts: int, top_k: int):
+ """
+ Initialize the top-k gating mechanism.
+
+ Args:
+ input_size (`int`):
+ Size of the input.
+ num_experts (`int`):
+ Number of experts.
+ top_k (`int`):
+ Number of top experts to select.
+ """
+ super().__init__()
+
+ self.num_experts = num_experts
+ self.input_size = input_size
+ self.top_k = top_k
+
+ self.layer = nn.Linear(input_size, num_experts, bias=False)
+
+ def forward(self, hidden_states):
+ # compute the top_k routing decision
+ logits = self.layer(hidden_states).float() # [batch_size x seq_len, num_experts]
+ top_k_logits, top_k_indices = logits.topk(self.top_k, dim=1) # [num_tokens, top_k]
+ top_k_gates = torch.softmax(top_k_logits, dim=1).type_as(hidden_states) # [num_tokens, top_k]
+
+ # compute number of input given to each expert
+ zeros = torch.zeros(
+ [top_k_gates.size(0), self.num_experts], dtype=top_k_gates.dtype, device=top_k_gates.device
+ ) # [num_tokens, num_experts]
+ gates = zeros.scatter(1, top_k_indices, 1) # [num_tokens, num_experts]
+ expert_size = gates.long().sum(0) # [num_experts,]
+ # (This cause torch.compile to fail with `torch._dynamo.exc.Unsupported: Backend compiler failed with a fake tensor exception at`)
+ # (and `DataDependentOutputException`)
+ expert_size = expert_size.tolist()
+
+ # sort and group input tokens according to expert assignment
+ top_k_experts = top_k_indices.flatten() # [num_tokens * top_k]
+ _, index_sorted_experts = top_k_experts.sort(0) # [num_tokens * top_k]
+ batch_index = index_sorted_experts.div(self.top_k, rounding_mode="trunc") # [num_tokens * top_k]
+
+ # gather the gate values for grouped input tokens
+ top_k_gates = top_k_gates.flatten() # [num_tokens * top_k]
+ batch_gates = top_k_gates[index_sorted_experts] # [num_tokens * top_k]
+
+ return index_sorted_experts, batch_index, batch_gates, expert_size, logits
+
+
+class JetMoeMoE(nn.Module):
+ """
+ A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+
+ def __init__(self, config: JetMoeConfig):
+ super(JetMoeMoE, self).__init__()
+
+ self.input_size = config.hidden_size
+ self.hidden_size = config.intermediate_size
+ self.activation = ACT2FN[config.activation_function]
+ self.bias = torch.nn.Parameter(torch.empty(self.input_size))
+ self.input_linear = JetMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
+ self.output_linear = JetMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
+
+ self.router = JetMoeTopKGating(
+ input_size=self.input_size,
+ num_experts=config.num_local_experts,
+ top_k=config.num_experts_per_tok,
+ )
+
+ def forward(self, layer_input):
+ """
+ Forward pass of the mixture of experts layer.
+
+ Args:
+ layer_input (Tensor):
+ Input tensor.
+
+ Returns:
+ Tensor:
+ Output tensor.
+ Tensor:
+ Router logits.
+ """
+ bsz, length, emb_size = layer_input.size()
+ layer_input = layer_input.reshape(-1, emb_size)
+ _, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input)
+
+ expert_inputs = layer_input[batch_index]
+ hidden_states = self.input_linear(expert_inputs, expert_size)
+ chunked_hidden_states = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
+ expert_outputs = self.output_linear(hidden_states, expert_size)
+
+ expert_outputs = expert_outputs * batch_gates[:, None]
+
+ zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device)
+ layer_output = zeros.index_add(0, batch_index, expert_outputs)
+ layer_output = layer_output.view(bsz, length, self.input_size)
+ layer_output = layer_output + self.bias
+ return layer_output, router_logits
+
+
+class JetMoeMoA(nn.Module):
+ """
+ A Sparsely gated mixture of attention layer with pairs of query- and output-projections as experts.
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+
+ def __init__(self, config: JetMoeConfig):
+ super(JetMoeMoA, self).__init__()
+
+ self.num_experts = config.num_local_experts
+ self.input_size = config.hidden_size
+ self.hidden_size = config.kv_channels * config.num_key_value_heads
+ self.top_k = config.num_experts_per_tok
+ self.bias = torch.nn.Parameter(torch.empty(self.input_size))
+
+ self.input_linear = JetMoeParallelExperts(self.num_experts, self.input_size, self.hidden_size)
+ self.output_linear = JetMoeParallelExperts(self.num_experts, self.hidden_size, self.input_size)
+
+ self.router = JetMoeTopKGating(
+ input_size=self.input_size,
+ num_experts=self.num_experts,
+ top_k=self.top_k,
+ )
+
+ def map(self, layer_input):
+ """
+ Map inputs to attention experts according to routing decision and compute query projection inside each experts.
+ """
+
+ # Compute gating topology
+ bsz, length, emb_size = layer_input.size()
+ layer_input = layer_input.reshape(-1, emb_size) # [bsz * length, emb_size]
+ index_sorted_experts, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input)
+ topo_info = (index_sorted_experts, batch_index, batch_gates, expert_size)
+
+ # Group inputs according to topology and compute query projection
+ expert_inputs = layer_input[batch_index] # [bsz * length * top_k, emb_size]
+ expert_outputs = self.input_linear(expert_inputs, expert_size) # [bsz * length * top_k, hidden_size]
+
+ # Ungroup queries back to original order
+ zeros = torch.zeros(
+ (bsz * length * self.top_k, self.hidden_size), dtype=expert_outputs.dtype, device=expert_outputs.device
+ )
+ layer_output = zeros.index_add(0, index_sorted_experts, expert_outputs)
+ layer_output = layer_output.view(bsz, length, self.top_k, -1) # [bsz, length, top_k, hidden_size]
+ return layer_output, router_logits, topo_info
+
+ def reduce(self, layer_input, topo_info):
+ """
+ Compute output projection inside each attention experts and merge the outputs of different experts.
+ """
+ bsz, length, k, hidden_size = layer_input.size()
+ layer_input = layer_input.reshape(-1, hidden_size) # [bsz * length * k, hidden_size]
+ index_sorted_experts, batch_index, batch_gates, expert_size = topo_info
+
+ # Group inputs according to topology and compute output projection
+ expert_inputs = layer_input[index_sorted_experts] # [bsz * length * top_k, hidden_size]
+ expert_outputs = self.output_linear(expert_inputs, expert_size) # [bsz * length * top_k, emb_size]
+
+ # Apply gates to attention expert outputs
+ expert_outputs = expert_outputs * batch_gates[:, None]
+
+ # Ungroup and merge outputs to original order
+ zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device)
+ layer_output = zeros.index_add(0, batch_index, expert_outputs)
+ layer_output = layer_output.view(bsz, length, self.input_size)
+ layer_output = layer_output + self.bias
+ return layer_output
+
+ def forward(self, layer_input):
+ raise NotImplementedError("This module doesn't support call and forward.")
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->JetMoe
+class JetMoeRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ JetMoeRMSNorm 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+# Copied from transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding with Gemma->JetMoe
+class JetMoeRotaryEmbedding(nn.Module):
+ def __init__(self, config: JetMoeConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ 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 JetMoeAttention(nn.Module):
+ """
+ Multi-headed attention from 'Attention Is All You Need' paper.
+ """
+
+ def __init__(self, config: JetMoeConfig, layer_idx: Optional[int] = None):
+ """
+ Initialize the JetMoeAttention module.
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ layer_idx:
+ Index of the layer in the model.
+ """
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.is_causal = True
+ 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.top_k = config.num_experts_per_tok
+ self.attention_dropout = config.attention_dropout
+ self.kv_projection_size = config.kv_channels * config.num_key_value_heads
+ self.num_key_value_heads = config.num_key_value_heads
+ self.num_heads = config.num_attention_heads
+ self.head_dim = config.kv_channels
+
+ self.experts = JetMoeMoA(config)
+
+ self.kv_proj = torch.nn.Linear(config.hidden_size, self.kv_projection_size * 2, bias=False)
+
+ self.rotary_emb = JetMoeRotaryEmbedding(config)
+
+ 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]]]:
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states, router_logits, topo_info = self.experts.map(hidden_states)
+ key_states, value_states = self.kv_proj(hidden_states).chunk(2, dim=-1)
+
+ 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)
+
+ 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)
+
+ # repeat k/v heads for top-k attention experts
+ key_states = key_states.repeat(1, self.top_k, 1, 1)
+ value_states = value_states.repeat(1, self.top_k, 1, 1)
+
+ 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.top_k, self.kv_projection_size)
+
+ attn_output = self.experts.reduce(attn_output, topo_info)
+ attn_output = attn_output.view(bsz, q_len, -1)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value, router_logits
+
+
+class JetMoeSdpaAttention(JetMoeAttention):
+ """
+ JetMoe attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+ `JetMoeAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+ SDPA API.
+ """
+
+ # Adapted from JetMoeAttention.forward
+ 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]], Optional[torch.Tensor]]:
+ if output_attentions:
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+ logger.warning_once(
+ "JetMoeModel is using JetMoeSdpaAttention, 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, router_logits, topo_info = self.experts.map(hidden_states)
+ key_states, value_states = self.kv_proj(hidden_states).chunk(2, dim=-1)
+
+ 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)
+
+ 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)
+
+ # repeat k/v heads for top-k attention experts
+ key_states = key_states.repeat(1, self.top_k, 1, 1)
+ value_states = value_states.repeat(1, self.top_k, 1, 1)
+
+ 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()
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ is_causal = True if causal_mask is None and q_len > 1 else False
+
+ 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=is_causal,
+ )
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+ attn_output = attn_output.reshape(bsz, q_len, self.top_k, self.kv_projection_size)
+
+ attn_output = self.experts.reduce(attn_output, topo_info)
+ attn_output = attn_output.view(bsz, q_len, -1)
+
+ return attn_output, None, past_key_value, router_logits
+
+
+class JetMoeFlashAttention2(JetMoeAttention):
+ 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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ def forward(
+ self,
+ hidden_states: Optional[torch.FloatTensor],
+ attention_mask: Optional[torch.FloatTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ output_attentions: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ ) -> Union[
+ Tuple[torch.Tensor, Tuple[torch.Tensor]],
+ Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
+ ]:
+ """
+ Forward pass of the JetMoeAttention module.
+
+ Args:
+ hidden_states (Optional[torch.FloatTensor]): Input hidden states.
+ attention_mask (Optional[torch.FloatTensor]): Attention mask.
+ layer_past (Optional[Tuple[torch.Tensor]]): Past layer state.
+ use_cache (Optional[bool]): Whether to use cached states.
+ output_attentions (Optional[bool]): Whether to output attention weights.
+ cache_position (Optional[torch.LongTensor]): Position of the cache.
+
+ Returns:
+ Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Optional[Tuple[...]]]: Tuple containing outputs.
+ """
+ output_attentions = False
+ bsz, q_len, hidden_size = hidden_states.size()
+
+ # calculate query, key, values
+ query_states, router_logits, topo_info = self.experts.map(hidden_states)
+ key_states, value_states = self.kv_proj(hidden_states).chunk(2, dim=-1)
+
+ 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)
+
+ 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)
+
+ # repeat k/v heads for top-k attention experts
+ key_states = key_states.repeat(1, self.top_k, 1, 1)
+ value_states = value_states.repeat(1, self.top_k, 1, 1)
+
+ # 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. (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 = self.kv_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 = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=dropout_rate,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ is_causal=self.is_causal,
+ ).to(input_dtype)
+
+ # output projection
+ attn_output = attn_output.reshape(bsz, q_len, self.top_k, self.kv_projection_size)
+ attn_output = self.experts.reduce(attn_output, topo_info)
+ attn_output = attn_output.view(bsz, q_len, hidden_size) # re-assemble all head outputs side by side
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value, router_logits
+
+
+JETMOE_ATTENTION_CLASSES = {
+ "eager": JetMoeAttention,
+ "flash_attention_2": JetMoeFlashAttention2,
+ "sdpa": JetMoeSdpaAttention,
+}
+
+
+class JetMoeBlock(nn.Module):
+ def __init__(self, config: JetMoeConfig, layer_idx: Optional[int] = None):
+ """
+ Initialize the JetMoeBlock module.
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+ super().__init__()
+ self.input_layernorm = JetMoeRMSNorm(config.hidden_size)
+ self.self_attention = JETMOE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+ self.post_attention_layernorm = JetMoeRMSNorm(config.hidden_size)
+
+ self.mlp = JetMoeMoE(config)
+
+ def forward(
+ self,
+ hidden_states: Optional[torch.FloatTensor],
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = False,
+ output_router_logits: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
+ # Self Attention
+ attn_output, self_attn_weights, present_key_value, attn_router_logits = self.self_attention(
+ hidden_states=self.input_layernorm(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,
+ )
+
+ hidden_states = hidden_states + attn_output
+ x_mlp, mlp_router_logits = self.mlp(self.post_attention_layernorm(hidden_states))
+ hidden_states = hidden_states + x_mlp
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights,)
+
+ if use_cache:
+ outputs += (present_key_value,)
+
+ if output_router_logits:
+ outputs += attn_router_logits, mlp_router_logits
+
+ return outputs
+
+
+class JetMoePreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = JetMoeConfig
+ base_model_prefix = "transformer"
+ supports_gradient_checkpointing = False
+ _no_split_modules = ["JetMoeBlock"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_cache_class = True
+
+ 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)
+ elif isinstance(module, JetMoeParallelExperts):
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, JetMoeMoA):
+ module.bias.data.zero_()
+ elif isinstance(module, JetMoeMoE):
+ module.bias.data.zero_()
+
+
+JETMOE_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 ([`JetMoeConfig`]): 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.
+"""
+
+JETMOE_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)
+ 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)
+ 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.
+ 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 JetMoe Model outputting raw hidden-states without any specific head on top.",
+ JETMOE_START_DOCSTRING,
+)
+class JetMoeModel(JetMoePreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`JetMoeBlock`]
+
+ Args:
+ config:
+ JetMoeConfig
+ """
+
+ def __init__(self, config: JetMoeConfig):
+ 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([JetMoeBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+ self._attn_implementation = config._attn_implementation
+ self.norm = JetMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel.get_input_embeddings
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel.set_input_embeddings
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(JETMOE_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[Union[Cache, 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,
+ output_router_logits: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ ) -> 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
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ 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)
+
+ # kept for BC (non `Cache` `past_key_values` inputs)
+ return_legacy_cache = False
+ if use_cache and not isinstance(past_key_values, Cache):
+ return_legacy_cache = True
+ if past_key_values is None:
+ past_key_values = DynamicCache()
+ else:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ logger.warning_once(
+ "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
+ "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
+ "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
+ )
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ 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)
+
+ if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
+ batch_size = inputs_embeds.shape[0]
+ is_padding_right = attention_mask[:, -1].sum().item() != batch_size
+ if is_padding_right:
+ raise ValueError(
+ "You are attempting to perform batched generation with padding_side='right'"
+ " this may lead to unexpected behaviour for Flash Attention version of JetMoe. Make sure to "
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
+ )
+ causal_mask = self._update_causal_mask(
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+ )
+
+ 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 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,
+ position_ids,
+ past_key_values,
+ causal_mask,
+ output_attentions,
+ output_router_logits,
+ use_cache,
+ use_reentrant=False,
+ )
+ 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,
+ output_router_logits=output_router_logits,
+ 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],)
+
+ if output_router_logits:
+ all_router_logits += (layer_outputs[-2], 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 return_legacy_cache:
+ next_cache = next_cache.to_legacy_cache()
+
+ 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,
+ )
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask
+ def _update_causal_mask(
+ self,
+ attention_mask: torch.Tensor,
+ input_tensor: torch.Tensor,
+ cache_position: torch.Tensor,
+ past_key_values: Cache,
+ output_attentions: bool = False,
+ ):
+ if self.config._attn_implementation == "flash_attention_2":
+ if attention_mask is not None and (attention_mask == 0.0).any():
+ return attention_mask
+ return None
+ if self.config._attn_implementation == "flex_attention":
+ if isinstance(attention_mask, torch.Tensor):
+ attention_mask = make_flex_block_causal_mask(attention_mask)
+ if isinstance(attention_mask, BlockMask):
+ return attention_mask
+
+ # 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. This feature is not compatible with static cache, as SDPA will fail
+ # to infer the attention mask.
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ using_static_cache = isinstance(past_key_values, StaticCache)
+
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
+ attention_mask,
+ inputs_embeds=input_tensor,
+ past_key_values_length=past_seen_tokens,
+ is_training=self.training,
+ ):
+ return None
+
+ dtype, device = input_tensor.dtype, input_tensor.device
+ sequence_length = input_tensor.shape[1]
+ if using_static_cache:
+ target_length = past_key_values.get_max_cache_shape()
+ else:
+ target_length = (
+ attention_mask.shape[-1]
+ if isinstance(attention_mask, torch.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask,
+ sequence_length=sequence_length,
+ target_length=target_length,
+ dtype=dtype,
+ device=device,
+ cache_position=cache_position,
+ batch_size=input_tensor.shape[0],
+ )
+
+ if (
+ self.config._attn_implementation == "sdpa"
+ and attention_mask is not None
+ and attention_mask.device.type in ["cuda", "xpu"]
+ and not output_attentions
+ ):
+ # 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
+ min_dtype = torch.finfo(dtype).min
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+ return causal_mask
+
+ @staticmethod
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel._prepare_4d_causal_attention_mask_with_cache_position
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ **kwargs,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
+ `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache,
+ to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ min_dtype = torch.finfo(dtype).min
+ 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(batch_size, 1, -1, -1)
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+ causal_mask.device
+ )
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+
+class JetMoeForCausalLM(JetMoePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = JetMoeModel(config)
+ self.vocab_size = config.vocab_size
+ self.aux_loss_coef = config.aux_loss_coef
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+ self.tie_word_embeddings = config.tie_word_embeddings
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
+ def get_decoder(self):
+ return self.model
+
+ @can_return_tuple
+ @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+ @add_start_docstrings_to_model_forward(JETMOE_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[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,
+ output_router_logits: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs,
+ ) -> MoeCausalLMOutputWithPast:
+ 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+ 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
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: MoeModelOutputWithPast = 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,
+ cache_position=cache_position,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ # Upcast to float if we need to compute the loss to avoid potential precision issues
+ logits = logits.float()
+ # Shift so that tokens < n predict n
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
+ shift_labels = shift_labels.view(-1)
+ # Ensure tensors are on the same device
+ shift_labels = shift_labels.to(shift_logits.device)
+ loss = self.loss_function(
+ shift_logits,
+ shift_labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ aux_loss = None
+ if output_router_logits:
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.num_experts,
+ self.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
+
+ 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,
+ )
+
+
+@add_start_docstrings(
+ """
+ The JetMoe Model transformer with a sequence classification head on top (linear layer).
+
+ [`JetMoeForSequenceClassification`] 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).
+ """,
+ JETMOE_START_DOCSTRING,
+)
+# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->JetMoe, LLAMA->JETMOE, BaseModelOutputWithPast->MoeModelOutputWithPast
+class JetMoeForSequenceClassification(JetMoePreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = JetMoeModel(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
+
+ @can_return_tuple
+ @add_start_docstrings_to_model_forward(JETMOE_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.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ ) -> 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).
+ """
+
+ transformer_outputs: MoeModelOutputWithPast = 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,
+ )
+ hidden_states = transformer_outputs.last_hidden_state
+ 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:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ 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), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+ 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,
+ )
+
+
+__all__ = ["JetMoeForCausalLM", "JetMoeModel", "JetMoePreTrainedModel", "JetMoeForSequenceClassification"]
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__init__.py b/lib/python3.12/site-packages/transformers/models/layoutlm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f079c33c7157898807c5f405086804cc1533ff9
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_layoutlm import *
+ from .modeling_layoutlm import *
+ from .modeling_tf_layoutlm import *
+ from .tokenization_layoutlm import *
+ from .tokenization_layoutlm_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6a45079dbfbe13d97b80251e7bd5a01fe81def18
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/configuration_layoutlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/configuration_layoutlm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2e9c3bba38b6fa6a84543506f22db6e1783695f3
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/configuration_layoutlm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_layoutlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_layoutlm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d9dac67d3204d745b8b164f4f18209131ae038ff
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_layoutlm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_tf_layoutlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_tf_layoutlm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7246e23684c8335c13af78871d48eda0fd513409
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/modeling_tf_layoutlm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a784a85ebd0b293dbd67d4e51926553a6e03ed19
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d6ae0c3f1af4fd6d4964914a2df69b4ec8aedfc1
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/layoutlm/__pycache__/tokenization_layoutlm_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/configuration_layoutlm.py b/lib/python3.12/site-packages/transformers/models/layoutlm/configuration_layoutlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..aebd25d5369088837a780721c2afaf3204f3004c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/configuration_layoutlm.py
@@ -0,0 +1,199 @@
+# coding=utf-8
+# Copyright 2010, The Microsoft Research Asia LayoutLM Team authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""LayoutLM model configuration"""
+
+from collections import OrderedDict
+from typing import Any, List, Mapping, Optional
+
+from ... import PretrainedConfig, PreTrainedTokenizer
+from ...onnx import OnnxConfig, PatchingSpec
+from ...utils import TensorType, is_torch_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class LayoutLMConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`LayoutLMModel`]. It is used to instantiate a
+ LayoutLM 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 LayoutLM
+ [microsoft/layoutlm-base-uncased](https://huggingface.co/microsoft/layoutlm-base-uncased) architecture.
+
+ Configuration objects inherit from [`BertConfig`] and can be used to control the model outputs. Read the
+ documentation from [`BertConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 30522):
+ Vocabulary size of the LayoutLM model. Defines the different tokens that can be represented by the
+ *inputs_ids* passed to the forward method of [`LayoutLMModel`].
+ 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 [`LayoutLMModel`].
+ 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):
+ The value used to pad input_ids.
+ 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`.
+ max_2d_position_embeddings (`int`, *optional*, defaults to 1024):
+ The maximum value that the 2D position embedding might ever used. Typically set this to something large
+ just in case (e.g., 1024).
+
+ Examples:
+
+ ```python
+ >>> from transformers import LayoutLMConfig, LayoutLMModel
+
+ >>> # Initializing a LayoutLM configuration
+ >>> configuration = LayoutLMConfig()
+
+ >>> # Initializing a model (with random weights) from the configuration
+ >>> model = LayoutLMModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "layoutlm"
+
+ 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",
+ use_cache=True,
+ max_2d_position_embeddings=1024,
+ **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.position_embedding_type = position_embedding_type
+ self.use_cache = use_cache
+ self.max_2d_position_embeddings = max_2d_position_embeddings
+
+
+class LayoutLMOnnxConfig(OnnxConfig):
+ def __init__(
+ self,
+ config: PretrainedConfig,
+ task: str = "default",
+ patching_specs: List[PatchingSpec] = None,
+ ):
+ super().__init__(config, task=task, patching_specs=patching_specs)
+ self.max_2d_positions = config.max_2d_position_embeddings - 1
+
+ @property
+ def inputs(self) -> Mapping[str, Mapping[int, str]]:
+ return OrderedDict(
+ [
+ ("input_ids", {0: "batch", 1: "sequence"}),
+ ("bbox", {0: "batch", 1: "sequence"}),
+ ("attention_mask", {0: "batch", 1: "sequence"}),
+ ("token_type_ids", {0: "batch", 1: "sequence"}),
+ ]
+ )
+
+ 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]:
+ """
+ Generate inputs to provide to the ONNX exporter for the specific framework
+
+ Args:
+ tokenizer: The tokenizer associated with this model configuration
+ batch_size: The batch size (int) to export the model for (-1 means dynamic axis)
+ seq_length: The sequence length (int) to export the model for (-1 means dynamic axis)
+ is_pair: Indicate if the input is a pair (sentence 1, sentence 2)
+ framework: The framework (optional) the tokenizer will generate tensor for
+
+ Returns:
+ Mapping[str, Tensor] holding the kwargs to provide to the model's forward function
+ """
+
+ input_dict = super().generate_dummy_inputs(
+ tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
+ )
+
+ # Generate a dummy bbox
+ box = [48, 84, 73, 128]
+
+ if not framework == TensorType.PYTORCH:
+ raise NotImplementedError("Exporting LayoutLM to ONNX is currently only supported for PyTorch.")
+
+ if not is_torch_available():
+ raise ValueError("Cannot generate dummy inputs without PyTorch installed.")
+ import torch
+
+ batch_size, seq_length = input_dict["input_ids"].shape
+ input_dict["bbox"] = torch.tensor([*[box] * seq_length]).tile(batch_size, 1, 1)
+ return input_dict
+
+
+__all__ = ["LayoutLMConfig", "LayoutLMOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_layoutlm.py b/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_layoutlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c31521a3f6d128bec84bd011e353e7818a2af6b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_layoutlm.py
@@ -0,0 +1,1387 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM 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 LayoutLM 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 (
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ MaskedLMOutput,
+ 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_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_layoutlm import LayoutLMConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LayoutLMConfig"
+_CHECKPOINT_FOR_DOC = "microsoft/layoutlm-base-uncased"
+
+
+LayoutLMLayerNorm = nn.LayerNorm
+
+
+class LayoutLMEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super(LayoutLMEmbeddings, self).__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.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+ self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+ self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+ self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids=None,
+ bbox=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]
+
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ 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=device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+
+ words_embeddings = inputs_embeds
+ position_embeddings = self.position_embeddings(position_ids)
+ try:
+ left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])
+ upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])
+ right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2])
+ lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3])
+ except IndexError as e:
+ raise IndexError("The `bbox`coordinate values should be within 0-1000 range.") from e
+
+ h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3] - bbox[:, :, 1])
+ w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2] - bbox[:, :, 0])
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+ embeddings = (
+ words_embeddings
+ + position_embeddings
+ + left_position_embeddings
+ + upper_position_embeddings
+ + right_position_embeddings
+ + lower_position_embeddings
+ + h_position_embeddings
+ + w_position_embeddings
+ + token_type_embeddings
+ )
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->LayoutLM
+class LayoutLMSelfAttention(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 LayoutLMModel 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->LayoutLM
+class LayoutLMSelfOutput(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
+
+
+LAYOUTLM_SELF_ATTENTION_CLASSES = {
+ "eager": LayoutLMSelfAttention,
+}
+
+
+# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->LayoutLM,BERT->LAYOUTLM
+class LayoutLMAttention(nn.Module):
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ self.self = LAYOUTLM_SELF_ATTENTION_CLASSES[config._attn_implementation](
+ config, position_embedding_type=position_embedding_type
+ )
+ self.output = LayoutLMSelfOutput(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 LayoutLMIntermediate(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->LayoutLM
+class LayoutLMOutput(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->LayoutLM
+class LayoutLMLayer(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 = LayoutLMAttention(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 = LayoutLMAttention(config, position_embedding_type="absolute")
+ self.intermediate = LayoutLMIntermediate(config)
+ self.output = LayoutLMOutput(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->LayoutLM
+class LayoutLMEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([LayoutLMLayer(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 LayoutLMPooler(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->LayoutLM
+class LayoutLMPredictionHeadTransform(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->LayoutLM
+class LayoutLMLMPredictionHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.transform = LayoutLMPredictionHeadTransform(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 _tie_weights(self):
+ 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->LayoutLM
+class LayoutLMOnlyMLMHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = LayoutLMLMPredictionHead(config)
+
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+ prediction_scores = self.predictions(sequence_output)
+ return prediction_scores
+
+
+class LayoutLMPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = LayoutLMConfig
+ base_model_prefix = "layoutlm"
+ 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, LayoutLMLayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+ elif isinstance(module, LayoutLMLMPredictionHead):
+ module.bias.data.zero_()
+
+
+LAYOUTLM_START_DOCSTRING = r"""
+ The LayoutLM model was proposed in [LayoutLM: Pre-training of Text and Layout for Document Image
+ Understanding](https://arxiv.org/abs/1912.13318) by Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang, Furu Wei and
+ Ming Zhou.
+
+ 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 ([`LayoutLMConfig`]): 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.
+"""
+
+LAYOUTLM_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)
+ bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*):
+ Bounding boxes of each input sequence tokens. Selected in the range `[0,
+ config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1)
+ format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1,
+ y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization.
+ 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 MASKED tokens.
+
+ [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 `(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*):
+ If set to `True`, the attentions tensors of all attention layers are returned. See `attentions` under
+ returned tensors for more detail.
+ output_hidden_states (`bool`, *optional*):
+ If set to `True`, the hidden states of all layers are returned. See `hidden_states` under returned tensors
+ for more detail.
+ return_dict (`bool`, *optional*):
+ If set to `True`, the model will return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare LayoutLM Model transformer outputting raw hidden-states without any specific head on top.",
+ LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMModel(LayoutLMPreTrainedModel):
+ def __init__(self, config):
+ super(LayoutLMModel, self).__init__(config)
+ self.config = config
+
+ self.embeddings = LayoutLMEmbeddings(config)
+ self.encoder = LayoutLMEncoder(config)
+ self.pooler = LayoutLMPooler(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(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=BaseModelOutputWithPoolingAndCrossAttentions, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ bbox: 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,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPoolingAndCrossAttentions]:
+ r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LayoutLMModel
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = LayoutLMModel.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="pt")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = torch.tensor([token_boxes])
+
+ >>> outputs = model(
+ ... input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids
+ ... )
+
+ >>> 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
+
+ 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)
+
+ if bbox is None:
+ bbox = torch.zeros(input_shape + (4,), dtype=torch.long, device=device)
+
+ extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+
+ extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
+ extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min
+
+ if head_mask is not None:
+ if head_mask.dim() == 1:
+ head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+ head_mask = head_mask.expand(self.config.num_hidden_layers, -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)
+ else:
+ head_mask = [None] * self.config.num_hidden_layers
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ bbox=bbox,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ 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]
+ pooled_output = self.pooler(sequence_output)
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_attentions,
+ )
+
+
+@add_start_docstrings("""LayoutLM Model with a `language modeling` head on top.""", LAYOUTLM_START_DOCSTRING)
+class LayoutLMForMaskedLM(LayoutLMPreTrainedModel):
+ _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.layoutlm = LayoutLMModel(config)
+ self.cls = LayoutLMOnlyMLMHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.layoutlm.embeddings.word_embeddings
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=MaskedLMOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ bbox: 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,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = 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]`
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LayoutLMForMaskedLM
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = LayoutLMForMaskedLM.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "[MASK]"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="pt")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = torch.tensor([token_boxes])
+
+ >>> labels = tokenizer("Hello world", return_tensors="pt")["input_ids"]
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=labels,
+ ... )
+
+ >>> loss = outputs.loss
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.layoutlm(
+ input_ids,
+ bbox,
+ 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.cls(sequence_output)
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ 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,
+ )
+
+
+@add_start_docstrings(
+ """
+ LayoutLM Model with a sequence classification head on top (a linear layer on top of the pooled output) e.g. for
+ document image classification tasks such as the [RVL-CDIP](https://www.cs.cmu.edu/~aharley/rvl-cdip/) dataset.
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMForSequenceClassification(LayoutLMPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.layoutlm = LayoutLMModel(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()
+
+ def get_input_embeddings(self):
+ return self.layoutlm.embeddings.word_embeddings
+
+ @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=SequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ bbox: 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).
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LayoutLMForSequenceClassification
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = LayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="pt")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = torch.tensor([token_boxes])
+ >>> sequence_label = torch.tensor([1])
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=sequence_label,
+ ... )
+
+ >>> loss = outputs.loss
+ >>> logits = outputs.logits
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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,
+ )
+
+ 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(
+ """
+ LayoutLM Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+ sequence labeling (information extraction) tasks such as the [FUNSD](https://guillaumejaume.github.io/FUNSD/)
+ dataset and the [SROIE](https://rrc.cvc.uab.es/?ch=13) dataset.
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMForTokenClassification(LayoutLMPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.layoutlm = LayoutLMModel(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()
+
+ def get_input_embeddings(self):
+ return self.layoutlm.embeddings.word_embeddings
+
+ @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TokenClassifierOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ bbox: 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]`.
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LayoutLMForTokenClassification
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="pt")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = torch.tensor([token_boxes])
+ >>> token_labels = torch.tensor([1, 1, 0, 0]).unsqueeze(0) # batch size of 1
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=token_labels,
+ ... )
+
+ >>> loss = outputs.loss
+ >>> logits = outputs.logits
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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[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(
+ """
+ LayoutLM Model with a span classification head on top for extractive question-answering tasks such as
+ [DocVQA](https://rrc.cvc.uab.es/?ch=17) (a linear layer on top of the final hidden-states output to compute `span
+ start logits` and `span end logits`).
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMForQuestionAnswering(LayoutLMPreTrainedModel):
+ def __init__(self, config, has_visual_segment_embedding=True):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.layoutlm = LayoutLMModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.layoutlm.embeddings.word_embeddings
+
+ @replace_return_docstrings(output_type=QuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ bbox: 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.
+
+ Returns:
+
+ Example:
+
+ In the example below, we prepare a question + context pair for the LayoutLM model. It will give us a prediction
+ of what it thinks the answer is (the span of the answer within the texts parsed from the image).
+
+ ```python
+ >>> from transformers import AutoTokenizer, LayoutLMForQuestionAnswering
+ >>> from datasets import load_dataset
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("impira/layoutlm-document-qa", add_prefix_space=True)
+ >>> model = LayoutLMForQuestionAnswering.from_pretrained("impira/layoutlm-document-qa", revision="1e3ebac")
+
+ >>> dataset = load_dataset("nielsr/funsd", split="train", trust_remote_code=True)
+ >>> example = dataset[0]
+ >>> question = "what's his name?"
+ >>> words = example["words"]
+ >>> boxes = example["bboxes"]
+
+ >>> encoding = tokenizer(
+ ... question.split(), words, is_split_into_words=True, return_token_type_ids=True, return_tensors="pt"
+ ... )
+ >>> bbox = []
+ >>> for i, s, w in zip(encoding.input_ids[0], encoding.sequence_ids(0), encoding.word_ids(0)):
+ ... if s == 1:
+ ... bbox.append(boxes[w])
+ ... elif i == tokenizer.sep_token_id:
+ ... bbox.append([1000] * 4)
+ ... else:
+ ... bbox.append([0] * 4)
+ >>> encoding["bbox"] = torch.tensor([bbox])
+
+ >>> word_ids = encoding.word_ids(0)
+ >>> outputs = model(**encoding)
+ >>> loss = outputs.loss
+ >>> start_scores = outputs.start_logits
+ >>> end_scores = outputs.end_logits
+ >>> start, end = word_ids[start_scores.argmax(-1)], word_ids[end_scores.argmax(-1)]
+ >>> print(" ".join(words[start : end + 1]))
+ M. Hamann P. Harper, P. Martinez
+ ```"""
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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,
+ )
+
+
+__all__ = [
+ "LayoutLMForMaskedLM",
+ "LayoutLMForSequenceClassification",
+ "LayoutLMForTokenClassification",
+ "LayoutLMForQuestionAnswering",
+ "LayoutLMModel",
+ "LayoutLMPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_tf_layoutlm.py b/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_tf_layoutlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..c17d0f1ec573d05027a71c32442f1ff04406ea63
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/modeling_tf_layoutlm.py
@@ -0,0 +1,1692 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM 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 2.0 LayoutLM model."""
+
+from __future__ import annotations
+
+import math
+import warnings
+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 (
+ TFBaseModelOutputWithPastAndCrossAttentions,
+ TFBaseModelOutputWithPoolingAndCrossAttentions,
+ TFMaskedLMOutput,
+ TFQuestionAnsweringModelOutput,
+ TFSequenceClassifierOutput,
+ TFTokenClassifierOutput,
+)
+from ...modeling_tf_utils import (
+ TFMaskedLanguageModelingLoss,
+ TFModelInputType,
+ 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_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_layoutlm import LayoutLMConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LayoutLMConfig"
+
+
+class TFLayoutLMEmbeddings(keras.layers.Layer):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config: LayoutLMConfig, **kwargs):
+ super().__init__(**kwargs)
+
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.max_position_embeddings = config.max_position_embeddings
+ self.max_2d_position_embeddings = config.max_2d_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.hidden_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.hidden_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.hidden_size],
+ initializer=get_initializer(self.initializer_range),
+ )
+
+ with tf.name_scope("x_position_embeddings"):
+ self.x_position_embeddings = self.add_weight(
+ name="embeddings",
+ shape=[self.max_2d_position_embeddings, self.hidden_size],
+ initializer=get_initializer(self.initializer_range),
+ )
+
+ with tf.name_scope("y_position_embeddings"):
+ self.y_position_embeddings = self.add_weight(
+ name="embeddings",
+ shape=[self.max_2d_position_embeddings, self.hidden_size],
+ initializer=get_initializer(self.initializer_range),
+ )
+
+ with tf.name_scope("h_position_embeddings"):
+ self.h_position_embeddings = self.add_weight(
+ name="embeddings",
+ shape=[self.max_2d_position_embeddings, self.hidden_size],
+ initializer=get_initializer(self.initializer_range),
+ )
+
+ with tf.name_scope("w_position_embeddings"):
+ self.w_position_embeddings = self.add_weight(
+ name="embeddings",
+ shape=[self.max_2d_position_embeddings, self.hidden_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.hidden_size])
+
+ def call(
+ self,
+ input_ids: Optional[tf.Tensor] = None,
+ bbox: Optional[tf.Tensor] = None,
+ position_ids: Optional[tf.Tensor] = None,
+ token_type_ids: Optional[tf.Tensor] = None,
+ inputs_embeds: Optional[tf.Tensor] = None,
+ training: bool = False,
+ ) -> tf.Tensor:
+ """
+ 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 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)
+
+ if position_ids is None:
+ position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
+
+ if bbox is None:
+ bbox = bbox = tf.fill(input_shape + [4], value=0)
+ try:
+ left_position_embeddings = tf.gather(self.x_position_embeddings, bbox[:, :, 0])
+ upper_position_embeddings = tf.gather(self.y_position_embeddings, bbox[:, :, 1])
+ right_position_embeddings = tf.gather(self.x_position_embeddings, bbox[:, :, 2])
+ lower_position_embeddings = tf.gather(self.y_position_embeddings, bbox[:, :, 3])
+ except IndexError as e:
+ raise IndexError("The `bbox`coordinate values should be within 0-1000 range.") from e
+ h_position_embeddings = tf.gather(self.h_position_embeddings, bbox[:, :, 3] - bbox[:, :, 1])
+ w_position_embeddings = tf.gather(self.w_position_embeddings, bbox[:, :, 2] - bbox[:, :, 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
+ + left_position_embeddings
+ + upper_position_embeddings
+ + right_position_embeddings
+ + lower_position_embeddings
+ + h_position_embeddings
+ + w_position_embeddings
+ )
+ final_embeddings = self.LayerNorm(inputs=final_embeddings)
+ final_embeddings = self.dropout(inputs=final_embeddings, training=training)
+
+ return final_embeddings
+
+
+# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention with Bert->LayoutLM
+class TFLayoutLMSelfAttention(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **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 TFLayoutLMModel 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->LayoutLM
+class TFLayoutLMSelfOutput(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **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->LayoutLM
+class TFLayoutLMAttention(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **kwargs):
+ super().__init__(**kwargs)
+
+ self.self_attention = TFLayoutLMSelfAttention(config, name="self")
+ self.dense_output = TFLayoutLMSelfOutput(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->LayoutLM
+class TFLayoutLMIntermediate(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **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->LayoutLM
+class TFLayoutLMOutput(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **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->LayoutLM
+class TFLayoutLMLayer(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **kwargs):
+ super().__init__(**kwargs)
+
+ self.attention = TFLayoutLMAttention(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 = TFLayoutLMAttention(config, name="crossattention")
+ self.intermediate = TFLayoutLMIntermediate(config, name="intermediate")
+ self.bert_output = TFLayoutLMOutput(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->LayoutLM
+class TFLayoutLMEncoder(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **kwargs):
+ super().__init__(**kwargs)
+ self.config = config
+ self.layer = [TFLayoutLMLayer(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->LayoutLM
+class TFLayoutLMPooler(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **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.bert.modeling_tf_bert.TFBertPredictionHeadTransform with Bert->LayoutLM
+class TFLayoutLMPredictionHeadTransform(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, **kwargs):
+ super().__init__(**kwargs)
+
+ self.dense = keras.layers.Dense(
+ units=config.hidden_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(inputs=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])
+
+
+# Copied from transformers.models.bert.modeling_tf_bert.TFBertLMPredictionHead with Bert->LayoutLM
+class TFLayoutLMLMPredictionHead(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, input_embeddings: keras.layers.Layer, **kwargs):
+ super().__init__(**kwargs)
+
+ self.config = config
+ self.hidden_size = config.hidden_size
+
+ self.transform = TFLayoutLMPredictionHeadTransform(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.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
+
+
+# Copied from transformers.models.bert.modeling_tf_bert.TFBertMLMHead with Bert->LayoutLM
+class TFLayoutLMMLMHead(keras.layers.Layer):
+ def __init__(self, config: LayoutLMConfig, input_embeddings: keras.layers.Layer, **kwargs):
+ super().__init__(**kwargs)
+
+ self.predictions = TFLayoutLMLMPredictionHead(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 TFLayoutLMMainLayer(keras.layers.Layer):
+ config_class = LayoutLMConfig
+
+ def __init__(self, config: LayoutLMConfig, add_pooling_layer: bool = True, **kwargs):
+ super().__init__(**kwargs)
+
+ self.config = config
+
+ self.embeddings = TFLayoutLMEmbeddings(config, name="embeddings")
+ self.encoder = TFLayoutLMEncoder(config, name="encoder")
+ self.pooler = TFLayoutLMPooler(config, name="pooler") if add_pooling_layer else None
+
+ 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,
+ bbox: np.ndarray | tf.Tensor | 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,
+ 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 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)
+ if bbox is None:
+ bbox = tf.fill(dims=input_shape + [4], value=0)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ bbox=bbox,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ 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.
+ 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=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)
+
+ # 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,
+ # Need to pass these required positional arguments to `Encoder`
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=False,
+ 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,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_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, "pooler", None) is not None:
+ with tf.name_scope(self.pooler.name):
+ self.pooler.build(None)
+
+
+class TFLayoutLMPreTrainedModel(TFPreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = LayoutLMConfig
+ base_model_prefix = "layoutlm"
+
+ @property
+ def input_signature(self):
+ signature = super().input_signature
+ signature["bbox"] = tf.TensorSpec(shape=(None, None, 4), dtype=tf.int32, name="bbox")
+ return signature
+
+
+LAYOUTLM_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 ([`LayoutLMConfig`]): 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.
+"""
+
+LAYOUTLM_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)
+ bbox (`Numpy array` or `tf.Tensor` of shape `({0}, 4)`, *optional*):
+ Bounding Boxes of each input sequence tokens. Selected in the range `[0, config.max_2d_position_embeddings-
+ 1]`.
+ 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.
+ 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.
+ 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 LayoutLM Model transformer outputting raw hidden-states without any specific head on top.",
+ LAYOUTLM_START_DOCSTRING,
+)
+class TFLayoutLMModel(TFLayoutLMPreTrainedModel):
+ def __init__(self, config: LayoutLMConfig, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.layoutlm = TFLayoutLMMainLayer(config, name="layoutlm")
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(
+ output_type=TFBaseModelOutputWithPoolingAndCrossAttentions, config_class=_CONFIG_FOR_DOC
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ bbox: np.ndarray | tf.Tensor | 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,
+ 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"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, TFLayoutLMModel
+ >>> import tensorflow as tf
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = TFLayoutLMModel.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="tf")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = tf.convert_to_tensor([token_boxes])
+
+ >>> outputs = model(
+ ... input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids
+ ... )
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ ```"""
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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, "layoutlm", None) is not None:
+ with tf.name_scope(self.layoutlm.name):
+ self.layoutlm.build(None)
+
+
+@add_start_docstrings("""LayoutLM Model with a `language modeling` head on top.""", LAYOUTLM_START_DOCSTRING)
+class TFLayoutLMForMaskedLM(TFLayoutLMPreTrainedModel, TFMaskedLanguageModelingLoss):
+ # 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"pooler",
+ r"cls.seq_relationship",
+ r"cls.predictions.decoder.weight",
+ r"nsp___cls",
+ ]
+
+ def __init__(self, config: LayoutLMConfig, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ if config.is_decoder:
+ logger.warning(
+ "If you want to use `TFLayoutLMForMaskedLM` make sure `config.is_decoder=False` for "
+ "bi-directional self-attention."
+ )
+
+ self.layoutlm = TFLayoutLMMainLayer(config, add_pooling_layer=True, name="layoutlm")
+ self.mlm = TFLayoutLMMLMHead(config, input_embeddings=self.layoutlm.embeddings, name="mlm___cls")
+
+ def get_lm_head(self) -> keras.layers.Layer:
+ return self.mlm.predictions
+
+ 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.mlm.name + "/" + self.mlm.predictions.name
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TFMaskedLMOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ bbox: np.ndarray | tf.Tensor | 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` 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]`
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, TFLayoutLMForMaskedLM
+ >>> import tensorflow as tf
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = TFLayoutLMForMaskedLM.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "[MASK]"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="tf")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = tf.convert_to_tensor([token_boxes])
+
+ >>> labels = tokenizer("Hello world", return_tensors="tf")["input_ids"]
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=labels,
+ ... )
+
+ >>> loss = outputs.loss
+ ```"""
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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]
+ 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, "layoutlm", None) is not None:
+ with tf.name_scope(self.layoutlm.name):
+ self.layoutlm.build(None)
+ if getattr(self, "mlm", None) is not None:
+ with tf.name_scope(self.mlm.name):
+ self.mlm.build(None)
+
+
+@add_start_docstrings(
+ """
+ LayoutLM Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+ pooled output) e.g. for GLUE tasks.
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class TFLayoutLMForSequenceClassification(TFLayoutLMPreTrainedModel, TFSequenceClassificationLoss):
+ # 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: LayoutLMConfig, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+
+ self.layoutlm = TFLayoutLMMainLayer(config, name="layoutlm")
+ 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(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TFSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ bbox: np.ndarray | tf.Tensor | 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` 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).
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, TFLayoutLMForSequenceClassification
+ >>> import tensorflow as tf
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = TFLayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="tf")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = tf.convert_to_tensor([token_boxes])
+ >>> sequence_label = tf.convert_to_tensor([1])
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=sequence_label,
+ ... )
+
+ >>> loss = outputs.loss
+ >>> logits = outputs.logits
+ ```"""
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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,
+ )
+ pooled_output = outputs[1]
+ pooled_output = self.dropout(inputs=pooled_output, training=training)
+ logits = self.classifier(inputs=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, "layoutlm", None) is not None:
+ with tf.name_scope(self.layoutlm.name):
+ self.layoutlm.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(
+ """
+ LayoutLM 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.
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class TFLayoutLMForTokenClassification(TFLayoutLMPreTrainedModel, TFTokenClassificationLoss):
+ # 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"pooler",
+ 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: LayoutLMConfig, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+
+ self.layoutlm = TFLayoutLMMainLayer(config, add_pooling_layer=True, name="layoutlm")
+ 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(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TFTokenClassifierOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ bbox: np.ndarray | tf.Tensor | 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` 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]`.
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> import tensorflow as tf
+ >>> from transformers import AutoTokenizer, TFLayoutLMForTokenClassification
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+ >>> model = TFLayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased")
+
+ >>> words = ["Hello", "world"]
+ >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+
+ >>> token_boxes = []
+ >>> for word, box in zip(words, normalized_word_boxes):
+ ... word_tokens = tokenizer.tokenize(word)
+ ... token_boxes.extend([box] * len(word_tokens))
+ >>> # add bounding boxes of cls + sep tokens
+ >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+ >>> encoding = tokenizer(" ".join(words), return_tensors="tf")
+ >>> input_ids = encoding["input_ids"]
+ >>> attention_mask = encoding["attention_mask"]
+ >>> token_type_ids = encoding["token_type_ids"]
+ >>> bbox = tf.convert_to_tensor([token_boxes])
+ >>> token_labels = tf.convert_to_tensor([1, 1, 0, 0])
+
+ >>> outputs = model(
+ ... input_ids=input_ids,
+ ... bbox=bbox,
+ ... attention_mask=attention_mask,
+ ... token_type_ids=token_type_ids,
+ ... labels=token_labels,
+ ... )
+
+ >>> loss = outputs.loss
+ >>> logits = outputs.logits
+ ```"""
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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(inputs=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[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, "layoutlm", None) is not None:
+ with tf.name_scope(self.layoutlm.name):
+ self.layoutlm.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(
+ """
+ LayoutLM Model with a span classification head on top for extractive question-answering tasks such as
+ [DocVQA](https://rrc.cvc.uab.es/?ch=17) (a linear layer on top of the final hidden-states output to compute `span
+ start logits` and `span end logits`).
+ """,
+ LAYOUTLM_START_DOCSTRING,
+)
+class TFLayoutLMForQuestionAnswering(TFLayoutLMPreTrainedModel, TFQuestionAnsweringLoss):
+ # 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"pooler",
+ r"mlm___cls",
+ r"nsp___cls",
+ r"cls.predictions",
+ r"cls.seq_relationship",
+ ]
+
+ def __init__(self, config: LayoutLMConfig, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.num_labels = config.num_labels
+
+ self.layoutlm = TFLayoutLMMainLayer(config, add_pooling_layer=True, name="layoutlm")
+ 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(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TFQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ bbox: np.ndarray | tf.Tensor | 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` 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.
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> import tensorflow as tf
+ >>> from transformers import AutoTokenizer, TFLayoutLMForQuestionAnswering
+ >>> from datasets import load_dataset
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("impira/layoutlm-document-qa", add_prefix_space=True)
+ >>> model = TFLayoutLMForQuestionAnswering.from_pretrained("impira/layoutlm-document-qa", revision="1e3ebac")
+
+ >>> dataset = load_dataset("nielsr/funsd", split="train", trust_remote_code=True)
+ >>> example = dataset[0]
+ >>> question = "what's his name?"
+ >>> words = example["words"]
+ >>> boxes = example["bboxes"]
+
+ >>> encoding = tokenizer(
+ ... question.split(), words, is_split_into_words=True, return_token_type_ids=True, return_tensors="tf"
+ ... )
+ >>> bbox = []
+ >>> for i, s, w in zip(encoding.input_ids[0], encoding.sequence_ids(0), encoding.word_ids(0)):
+ ... if s == 1:
+ ... bbox.append(boxes[w])
+ ... elif i == tokenizer.sep_token_id:
+ ... bbox.append([1000] * 4)
+ ... else:
+ ... bbox.append([0] * 4)
+ >>> encoding["bbox"] = tf.convert_to_tensor([bbox])
+
+ >>> word_ids = encoding.word_ids(0)
+ >>> outputs = model(**encoding)
+ >>> loss = outputs.loss
+ >>> start_scores = outputs.start_logits
+ >>> end_scores = outputs.end_logits
+ >>> start, end = word_ids[tf.math.argmax(start_scores, -1)[0]], word_ids[tf.math.argmax(end_scores, -1)[0]]
+ >>> print(" ".join(words[start : end + 1]))
+ M. Hamann P. Harper, P. Martinez
+ ```"""
+
+ outputs = self.layoutlm(
+ input_ids=input_ids,
+ bbox=bbox,
+ 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(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, "layoutlm", None) is not None:
+ with tf.name_scope(self.layoutlm.name):
+ self.layoutlm.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])
+
+
+__all__ = [
+ "TFLayoutLMForMaskedLM",
+ "TFLayoutLMForSequenceClassification",
+ "TFLayoutLMForTokenClassification",
+ "TFLayoutLMForQuestionAnswering",
+ "TFLayoutLMMainLayer",
+ "TFLayoutLMModel",
+ "TFLayoutLMPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm.py b/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0b61c93ac1aa87dfce99f9150e207186e435e7c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm.py
@@ -0,0 +1,512 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT 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 LayoutLM."""
+
+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->LayoutLM,BERT->LayoutLM
+class LayoutLMTokenizer(PreTrainedTokenizer):
+ r"""
+ Construct a LayoutLM 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 LayoutLM).
+ clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`):
+ Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
+ extra spaces.
+ """
+
+ 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,
+ clean_up_tokenization_spaces=True,
+ **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 = LayoutLMTokenizer.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,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **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 LayoutLM 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 LayoutLM 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:
+ """
+ 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:
+ """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
+
+
+__all__ = ["LayoutLMTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm_fast.py b/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..1df94cb625a368e9c39d4b8ecdd17c9f5d210c73
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/layoutlm/tokenization_layoutlm_fast.py
@@ -0,0 +1,176 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT 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 LayoutLM."""
+
+import json
+from typing import List, Optional, Tuple
+
+from tokenizers import normalizers
+
+from ...tokenization_utils_fast import PreTrainedTokenizerFast
+from ...utils import logging
+from .tokenization_layoutlm import LayoutLMTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
+
+
+# Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast with Bert->LayoutLM,BERT->LayoutLM
+class LayoutLMTokenizerFast(PreTrainedTokenizerFast):
+ r"""
+ Construct a "fast" LayoutLM 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 LayoutLM).
+ wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
+ The prefix for subwords.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ slow_tokenizer_class = LayoutLMTokenizer
+
+ 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 LayoutLM 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 LayoutLM 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)
+
+
+__all__ = ["LayoutLMTokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/__init__.py b/lib/python3.12/site-packages/transformers/models/longformer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..87f53105424b76e5c18bd740ecfdd37a5b29d0d4
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_longformer import *
+ from .modeling_longformer import *
+ from .modeling_tf_longformer import *
+ from .tokenization_longformer import *
+ from .tokenization_longformer_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9e7f17b979bf50949a75b3ccd6494c63d656d6fe
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/configuration_longformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/configuration_longformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..361cb3f9dabea6f68b49fcc3b849b6aa0dbc3477
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/configuration_longformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6f4200b5c189aaa77b699e157897484936805bf8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..470984854833a5bf47f450a7ef65034b4723cbc7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/longformer/__pycache__/tokenization_longformer_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/configuration_longformer.py b/lib/python3.12/site-packages/transformers/models/longformer/configuration_longformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a4d698471750b88df18e4c97be9bae33b19fa50
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/configuration_longformer.py
@@ -0,0 +1,204 @@
+# coding=utf-8
+# Copyright 2020 The Allen Institute for AI team 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.
+"""Longformer configuration"""
+
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union
+
+from ...configuration_utils import PretrainedConfig
+from ...onnx import OnnxConfig
+from ...utils import TensorType, logging
+
+
+if TYPE_CHECKING:
+ from ...onnx.config import PatchingSpec
+ from ...tokenization_utils_base import PreTrainedTokenizerBase
+
+
+logger = logging.get_logger(__name__)
+
+
+class LongformerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`LongformerModel`] or a [`TFLongformerModel`]. It
+ is used to instantiate a Longformer model according to the specified arguments, defining the model architecture.
+
+ This is the configuration class to store the configuration of a [`LongformerModel`]. It is used to instantiate an
+ Longformer 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 LongFormer
+ [allenai/longformer-base-4096](https://huggingface.co/allenai/longformer-base-4096) architecture with a sequence
+ length 4,096.
+
+ 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 Longformer model. Defines the number of different tokens that can be represented by
+ the `inputs_ids` passed when calling [`LongformerModel`] or [`TFLongformerModel`].
+ 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 [`LongformerModel`] or
+ [`TFLongformerModel`].
+ 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.
+ attention_window (`int` or `List[int]`, *optional*, defaults to 512):
+ Size of an attention window around each token. If an `int`, use the same size for all layers. To specify a
+ different window size for each layer, use a `List[int]` where `len(attention_window) == num_hidden_layers`.
+
+ Example:
+
+ ```python
+ >>> from transformers import LongformerConfig, LongformerModel
+
+ >>> # Initializing a Longformer configuration
+ >>> configuration = LongformerConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = LongformerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "longformer"
+
+ def __init__(
+ self,
+ attention_window: Union[List[int], int] = 512,
+ sep_token_id: int = 2,
+ pad_token_id: int = 1,
+ bos_token_id: int = 0,
+ eos_token_id: int = 2,
+ vocab_size: int = 30522,
+ 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 = 512,
+ type_vocab_size: int = 2,
+ initializer_range: float = 0.02,
+ layer_norm_eps: float = 1e-12,
+ onnx_export: bool = False,
+ **kwargs,
+ ):
+ """Constructs LongformerConfig."""
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+ self.attention_window = attention_window
+ self.sep_token_id = sep_token_id
+ self.bos_token_id = bos_token_id
+ self.eos_token_id = eos_token_id
+ 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.onnx_export = onnx_export
+
+
+class LongformerOnnxConfig(OnnxConfig):
+ def __init__(self, config: "PretrainedConfig", task: str = "default", patching_specs: "List[PatchingSpec]" = None):
+ super().__init__(config, task, patching_specs)
+ config.onnx_export = True
+
+ @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),
+ ("global_attention_mask", dynamic_axis),
+ ]
+ )
+
+ @property
+ def outputs(self) -> Mapping[str, Mapping[int, str]]:
+ outputs = super().outputs
+ if self.task == "default":
+ outputs["pooler_output"] = {0: "batch"}
+ return outputs
+
+ @property
+ def atol_for_validation(self) -> float:
+ """
+ What absolute tolerance value to use during model conversion validation.
+
+ Returns:
+ Float absolute tolerance value.
+ """
+ return 1e-4
+
+ @property
+ def default_onnx_opset(self) -> int:
+ # needs to be >= 14 to support tril operator
+ return max(super().default_onnx_opset, 14)
+
+ def generate_dummy_inputs(
+ self,
+ tokenizer: "PreTrainedTokenizerBase",
+ batch_size: int = -1,
+ seq_length: int = -1,
+ is_pair: bool = False,
+ framework: Optional[TensorType] = None,
+ ) -> Mapping[str, Any]:
+ inputs = super().generate_dummy_inputs(
+ preprocessor=tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
+ )
+ import torch
+
+ # for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
+ # makes the export fail randomly
+ inputs["global_attention_mask"] = torch.zeros_like(inputs["input_ids"])
+ # make every second token global
+ inputs["global_attention_mask"][:, ::2] = 1
+
+ return inputs
+
+
+__all__ = ["LongformerConfig", "LongformerOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/modeling_longformer.py b/lib/python3.12/site-packages/transformers/models/longformer/modeling_longformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..9eefa02a6687552c5eba8ac55ae07f58f7c6a86b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/modeling_longformer.py
@@ -0,0 +1,2335 @@
+# coding=utf-8
+# Copyright 2020 The Allen Institute for AI team 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 Longformer 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, gelu
+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_longformer import LongformerConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "allenai/longformer-base-4096"
+_CONFIG_FOR_DOC = "LongformerConfig"
+
+
+@dataclass
+class LongformerBaseModelOutput(ModelOutput):
+ """
+ Base class for Longformer's outputs, with potential hidden states, local and global 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 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ last_hidden_state: torch.FloatTensor
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerBaseModelOutputWithPooling(ModelOutput):
+ """
+ Base class for Longformer'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)`):
+ Last layer hidden-state of the first token of the sequence (classification token) further processed by 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 + 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ last_hidden_state: torch.FloatTensor
+ pooler_output: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerMaskedLMOutput(ModelOutput):
+ """
+ Base class for masked language models outputs.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Masked language modeling (MLM) 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).
+ 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerQuestionAnsweringModelOutput(ModelOutput):
+ """
+ Base class for outputs of question answering Longformer models.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+ start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Span-start scores (before SoftMax).
+ end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Span-end 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 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ start_logits: Optional[torch.FloatTensor] = None
+ end_logits: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerSequenceClassifierOutput(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).
+ 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerMultipleChoiceModelOutput(ModelOutput):
+ """
+ Base class for outputs of multiple choice Longformer models.
+
+ Args:
+ loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided):
+ Classification loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`):
+ *num_choices* is the second dimension of the input tensors. (see *input_ids* above).
+
+ Classification 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 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class LongformerTokenClassifierOutput(ModelOutput):
+ """
+ Base class for outputs of token classification models.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :
+ Classification loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):
+ Classification 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 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`,
+ where `x` is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+def _get_question_end_index(input_ids, sep_token_id):
+ """
+ Computes the index of the first occurrence of `sep_token_id`.
+ """
+
+ sep_token_indices = (input_ids == sep_token_id).nonzero()
+ batch_size = input_ids.shape[0]
+
+ assert sep_token_indices.shape[1] == 2, "`input_ids` should have two dimensions"
+ assert sep_token_indices.shape[0] == 3 * batch_size, (
+ f"There should be exactly three separator tokens: {sep_token_id} in every sample for questions answering. You"
+ " might also consider to set `global_attention_mask` manually in the forward function to avoid this error."
+ )
+ return sep_token_indices.view(batch_size, 3, 2)[:, 0, 1]
+
+
+def _compute_global_attention_mask(input_ids, sep_token_id, before_sep_token=True):
+ """
+ Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
+ True` else after `sep_token_id`.
+ """
+ question_end_index = _get_question_end_index(input_ids, sep_token_id)
+ question_end_index = question_end_index.unsqueeze(dim=1) # size: batch_size x 1
+ # bool attention mask with True in locations of global attention
+ attention_mask = torch.arange(input_ids.shape[1], device=input_ids.device)
+ if before_sep_token is True:
+ attention_mask = (attention_mask.expand_as(input_ids) < question_end_index).to(torch.bool)
+ else:
+ # last token is separation token and should not be counted and in the middle are two separation tokens
+ attention_mask = (attention_mask.expand_as(input_ids) > (question_end_index + 1)).to(torch.bool) * (
+ attention_mask.expand_as(input_ids) < input_ids.shape[-1]
+ ).to(torch.bool)
+
+ return attention_mask
+
+
+def create_position_ids_from_input_ids(input_ids, padding_idx):
+ """
+ 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) * mask
+ return incremental_indices.long() + padding_idx
+
+
+class LongformerEmbeddings(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)
+ 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)
+
+ 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):
+ 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).to(input_ids.device)
+ 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]
+
+ if token_type_ids is None:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=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
+
+ 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 inputs_embeds:
+
+ 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 LongformerSelfAttention(nn.Module):
+ def __init__(self, config, layer_id):
+ 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_heads = config.num_attention_heads
+ self.head_dim = int(config.hidden_size / config.num_attention_heads)
+ self.embed_dim = config.hidden_size
+
+ self.query = nn.Linear(config.hidden_size, self.embed_dim)
+ self.key = nn.Linear(config.hidden_size, self.embed_dim)
+ self.value = nn.Linear(config.hidden_size, self.embed_dim)
+
+ # separate projection layers for tokens with global attention
+ self.query_global = nn.Linear(config.hidden_size, self.embed_dim)
+ self.key_global = nn.Linear(config.hidden_size, self.embed_dim)
+ self.value_global = nn.Linear(config.hidden_size, self.embed_dim)
+
+ self.dropout = config.attention_probs_dropout_prob
+
+ self.layer_id = layer_id
+ attention_window = config.attention_window[self.layer_id]
+ assert attention_window % 2 == 0, (
+ f"`attention_window` for layer {self.layer_id} has to be an even value. Given {attention_window}"
+ )
+ assert attention_window > 0, (
+ f"`attention_window` for layer {self.layer_id} has to be positive. Given {attention_window}"
+ )
+
+ self.one_sided_attn_window_size = attention_window // 2
+
+ self.config = config
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ layer_head_mask=None,
+ is_index_masked=None,
+ is_index_global_attn=None,
+ is_global_attn=None,
+ output_attentions=False,
+ ):
+ """
+ [`LongformerSelfAttention`] expects *len(hidden_states)* to be multiple of *attention_window*. Padding to
+ *attention_window* happens in [`LongformerModel.forward`] to avoid redoing the padding on each layer.
+
+ The *attention_mask* is changed in [`LongformerModel.forward`] from 0, 1, 2 to:
+
+ - -10000: no attention
+ - 0: local attention
+ - +10000: global attention
+ """
+ hidden_states = hidden_states.transpose(0, 1)
+
+ # project hidden states
+ query_vectors = self.query(hidden_states)
+ key_vectors = self.key(hidden_states)
+ value_vectors = self.value(hidden_states)
+
+ seq_len, batch_size, embed_dim = hidden_states.size()
+ assert embed_dim == self.embed_dim, (
+ f"hidden_states should have embed_dim = {self.embed_dim}, but has {embed_dim}"
+ )
+
+ # normalize query
+ query_vectors /= math.sqrt(self.head_dim)
+
+ query_vectors = query_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1)
+ key_vectors = key_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1)
+
+ attn_scores = self._sliding_chunks_query_key_matmul(
+ query_vectors, key_vectors, self.one_sided_attn_window_size
+ )
+
+ # values to pad for attention probs
+ remove_from_windowed_attention_mask = (attention_mask != 0)[:, :, None, None]
+
+ # cast to fp32/fp16 then replace 1's with -inf
+ float_mask = remove_from_windowed_attention_mask.type_as(query_vectors).masked_fill(
+ remove_from_windowed_attention_mask, torch.finfo(query_vectors.dtype).min
+ )
+ # diagonal mask with zeros everywhere and -inf inplace of padding
+ diagonal_mask = self._sliding_chunks_query_key_matmul(
+ float_mask.new_ones(size=float_mask.size()), float_mask, self.one_sided_attn_window_size
+ )
+
+ # pad local attention probs
+ attn_scores += diagonal_mask
+
+ assert list(attn_scores.size()) == [
+ batch_size,
+ seq_len,
+ self.num_heads,
+ self.one_sided_attn_window_size * 2 + 1,
+ ], (
+ f"local_attn_probs should be of size ({batch_size}, {seq_len}, {self.num_heads},"
+ f" {self.one_sided_attn_window_size * 2 + 1}), but is of size {attn_scores.size()}"
+ )
+
+ # compute local attention probs from global attention keys and contact over window dim
+ if is_global_attn:
+ # compute global attn indices required through out forward fn
+ (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ ) = self._get_global_attn_indices(is_index_global_attn)
+ # calculate global attn probs from global key
+
+ global_key_attn_scores = self._concat_with_global_key_attn_probs(
+ query_vectors=query_vectors,
+ key_vectors=key_vectors,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ )
+ # concat to local_attn_probs
+ # (batch_size, seq_len, num_heads, extra attention count + 2*window+1)
+ attn_scores = torch.cat((global_key_attn_scores, attn_scores), dim=-1)
+
+ # free memory
+ del global_key_attn_scores
+
+ attn_probs = nn.functional.softmax(
+ attn_scores, dim=-1, dtype=torch.float32
+ ) # use fp32 for numerical stability
+
+ 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_probs = layer_head_mask.view(1, 1, -1, 1) * attn_probs
+
+ # softmax sometimes inserts NaN if all positions are masked, replace them with 0
+ attn_probs = torch.masked_fill(attn_probs, is_index_masked[:, :, None, None], 0.0)
+ attn_probs = attn_probs.type_as(attn_scores)
+
+ # free memory
+ del attn_scores
+
+ # apply dropout
+ attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training)
+
+ value_vectors = value_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1)
+
+ # compute local attention output with global attention value and add
+ if is_global_attn:
+ # compute sum of global and local attn
+ attn_output = self._compute_attn_output_with_global_indices(
+ value_vectors=value_vectors,
+ attn_probs=attn_probs,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ )
+ else:
+ # compute local attn only
+ attn_output = self._sliding_chunks_matmul_attn_probs_value(
+ attn_probs, value_vectors, self.one_sided_attn_window_size
+ )
+
+ assert attn_output.size() == (batch_size, seq_len, self.num_heads, self.head_dim), "Unexpected size"
+ attn_output = attn_output.transpose(0, 1).reshape(seq_len, batch_size, embed_dim).contiguous()
+
+ # compute value for global attention and overwrite to attention output
+ # TODO: remove the redundant computation
+ if is_global_attn:
+ global_attn_output, global_attn_probs = self._compute_global_attn_output_from_hidden(
+ hidden_states=hidden_states,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ layer_head_mask=layer_head_mask,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ is_index_masked=is_index_masked,
+ )
+
+ # get only non zero global attn output
+ nonzero_global_attn_output = global_attn_output[
+ is_local_index_global_attn_nonzero[0], :, is_local_index_global_attn_nonzero[1]
+ ]
+
+ # overwrite values with global attention
+ attn_output[is_index_global_attn_nonzero[::-1]] = nonzero_global_attn_output.view(
+ len(is_local_index_global_attn_nonzero[0]), -1
+ )
+ # The attention weights for tokens with global attention are
+ # just filler values, they were never used to compute the output.
+ # Fill with 0 now, the correct values are in 'global_attn_probs'.
+ attn_probs[is_index_global_attn_nonzero] = 0
+
+ outputs = (attn_output.transpose(0, 1),)
+
+ if output_attentions:
+ outputs += (attn_probs,)
+
+ return outputs + (global_attn_probs,) if (is_global_attn and output_attentions) else outputs
+
+ @staticmethod
+ def _pad_and_transpose_last_two_dims(hidden_states_padded, padding):
+ """pads rows and then flips rows and columns"""
+ hidden_states_padded = nn.functional.pad(
+ hidden_states_padded, padding
+ ) # padding value is not important because it will be overwritten
+ hidden_states_padded = hidden_states_padded.view(
+ *hidden_states_padded.size()[:-2], hidden_states_padded.size(-1), hidden_states_padded.size(-2)
+ )
+ return hidden_states_padded
+
+ @staticmethod
+ def _pad_and_diagonalize(chunked_hidden_states):
+ """
+ shift every row 1 step right, converting columns into diagonals.
+
+ Example:
+
+ ```python
+ chunked_hidden_states: [
+ 0.4983,
+ 2.6918,
+ -0.0071,
+ 1.0492,
+ -1.8348,
+ 0.7672,
+ 0.2986,
+ 0.0285,
+ -0.7584,
+ 0.4206,
+ -0.0405,
+ 0.1599,
+ 2.0514,
+ -1.1600,
+ 0.5372,
+ 0.2629,
+ ]
+ window_overlap = num_rows = 4
+ ```
+
+ (pad & diagonalize) => [ 0.4983, 2.6918, -0.0071, 1.0492, 0.0000, 0.0000, 0.0000
+ 0.0000, -1.8348, 0.7672, 0.2986, 0.0285, 0.0000, 0.0000 0.0000, 0.0000, -0.7584, 0.4206,
+ -0.0405, 0.1599, 0.0000 0.0000, 0.0000, 0.0000, 2.0514, -1.1600, 0.5372, 0.2629 ]
+ """
+ total_num_heads, num_chunks, window_overlap, hidden_dim = chunked_hidden_states.size()
+ chunked_hidden_states = nn.functional.pad(
+ chunked_hidden_states, (0, window_overlap + 1)
+ ) # total_num_heads x num_chunks x window_overlap x (hidden_dim+window_overlap+1). Padding value is not important because it'll be overwritten
+ chunked_hidden_states = chunked_hidden_states.view(
+ total_num_heads, num_chunks, -1
+ ) # total_num_heads x num_chunks x window_overlap*window_overlap+window_overlap
+ chunked_hidden_states = chunked_hidden_states[
+ :, :, :-window_overlap
+ ] # total_num_heads x num_chunks x window_overlap*window_overlap
+ chunked_hidden_states = chunked_hidden_states.view(
+ total_num_heads, num_chunks, window_overlap, window_overlap + hidden_dim
+ )
+ chunked_hidden_states = chunked_hidden_states[:, :, :, :-1]
+ return chunked_hidden_states
+
+ @staticmethod
+ def _chunk(hidden_states, window_overlap, onnx_export: bool = False):
+ """convert into overlapping chunks. Chunk size = 2w, overlap size = w"""
+ if not onnx_export:
+ # non-overlapping chunks of size = 2w
+ hidden_states = hidden_states.view(
+ hidden_states.size(0),
+ torch.div(hidden_states.size(1), (window_overlap * 2), rounding_mode="trunc"),
+ window_overlap * 2,
+ hidden_states.size(2),
+ )
+ # use `as_strided` to make the chunks overlap with an overlap size = window_overlap
+ chunk_size = list(hidden_states.size())
+ chunk_size[1] = chunk_size[1] * 2 - 1
+
+ chunk_stride = list(hidden_states.stride())
+ chunk_stride[1] = chunk_stride[1] // 2
+ return hidden_states.as_strided(size=chunk_size, stride=chunk_stride)
+
+ # When exporting to ONNX, use this separate logic
+ # have to use slow implementation since as_strided, unfold and 2d-tensor indexing aren't supported (yet) in ONNX export
+
+ # TODO replace this with
+ # > return hidden_states.unfold(dimension=1, size=window_overlap * 2, step=window_overlap).transpose(2, 3)
+ # once `unfold` is supported
+ # the case hidden_states.size(1) == window_overlap * 2 can also simply return hidden_states.unsqueeze(1), but that's control flow
+
+ chunk_size = [
+ hidden_states.size(0),
+ torch.div(hidden_states.size(1), window_overlap, rounding_mode="trunc") - 1,
+ window_overlap * 2,
+ hidden_states.size(2),
+ ]
+
+ overlapping_chunks = torch.empty(chunk_size, device=hidden_states.device)
+ for chunk in range(chunk_size[1]):
+ overlapping_chunks[:, chunk, :, :] = hidden_states[
+ :, chunk * window_overlap : chunk * window_overlap + 2 * window_overlap, :
+ ]
+ return overlapping_chunks
+
+ @staticmethod
+ def _mask_invalid_locations(input_tensor, affected_seq_len) -> torch.Tensor:
+ beginning_mask_2d = input_tensor.new_ones(affected_seq_len, affected_seq_len + 1).tril().flip(dims=[0])
+ beginning_mask = beginning_mask_2d[None, :, None, :]
+ ending_mask = beginning_mask.flip(dims=(1, 3))
+ beginning_input = input_tensor[:, :affected_seq_len, :, : affected_seq_len + 1]
+ beginning_mask = beginning_mask.expand(beginning_input.size())
+ input_tensor[:, :affected_seq_len, :, : affected_seq_len + 1] = torch.full_like(
+ beginning_input, -float("inf")
+ ).where(beginning_mask.bool(), beginning_input)
+ ending_input = input_tensor[:, -affected_seq_len:, :, -(affected_seq_len + 1) :]
+ ending_mask = ending_mask.expand(ending_input.size())
+ input_tensor[:, -affected_seq_len:, :, -(affected_seq_len + 1) :] = torch.full_like(
+ ending_input, -float("inf")
+ ).where(ending_mask.bool(), ending_input)
+
+ def _sliding_chunks_query_key_matmul(self, query: torch.Tensor, key: torch.Tensor, window_overlap: int):
+ """
+ Matrix multiplication of query and key tensors using with a sliding window attention pattern. This
+ implementation splits the input into overlapping chunks of size 2w (e.g. 512 for pretrained Longformer) with an
+ overlap of size window_overlap
+ """
+ batch_size, seq_len, num_heads, head_dim = query.size()
+ assert seq_len % (window_overlap * 2) == 0, (
+ f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}"
+ )
+ assert query.size() == key.size()
+
+ chunks_count = torch.div(seq_len, window_overlap, rounding_mode="trunc") - 1
+
+ # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size window_overlap * 2
+ query = query.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+ key = key.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+
+ query = self._chunk(query, window_overlap, getattr(self.config, "onnx_export", False))
+ key = self._chunk(key, window_overlap, getattr(self.config, "onnx_export", False))
+
+ # matrix multiplication
+ # bcxd: batch_size * num_heads x chunks x 2window_overlap x head_dim
+ # bcyd: batch_size * num_heads x chunks x 2window_overlap x head_dim
+ # bcxy: batch_size * num_heads x chunks x 2window_overlap x 2window_overlap
+ diagonal_chunked_attention_scores = torch.einsum("bcxd,bcyd->bcxy", (query, key)) # multiply
+
+ # convert diagonals into columns
+ diagonal_chunked_attention_scores = self._pad_and_transpose_last_two_dims(
+ diagonal_chunked_attention_scores, padding=(0, 0, 0, 1)
+ )
+
+ # allocate space for the overall attention matrix where the chunks are combined. The last dimension
+ # has (window_overlap * 2 + 1) columns. The first (window_overlap) columns are the window_overlap lower triangles (attention from a word to
+ # window_overlap previous words). The following column is attention score from each word to itself, then
+ # followed by window_overlap columns for the upper triangle.
+
+ diagonal_attention_scores = diagonal_chunked_attention_scores.new_zeros(
+ (batch_size * num_heads, chunks_count + 1, window_overlap, window_overlap * 2 + 1)
+ )
+
+ # copy parts from diagonal_chunked_attention_scores into the combined matrix of attentions
+ # - copying the main diagonal and the upper triangle
+ diagonal_attention_scores[:, :-1, :, window_overlap:] = diagonal_chunked_attention_scores[
+ :, :, :window_overlap, : window_overlap + 1
+ ]
+ diagonal_attention_scores[:, -1, :, window_overlap:] = diagonal_chunked_attention_scores[
+ :, -1, window_overlap:, : window_overlap + 1
+ ]
+ # - copying the lower triangle
+ diagonal_attention_scores[:, 1:, :, :window_overlap] = diagonal_chunked_attention_scores[
+ :, :, -(window_overlap + 1) : -1, window_overlap + 1 :
+ ]
+
+ diagonal_attention_scores[:, 0, 1:window_overlap, 1:window_overlap] = diagonal_chunked_attention_scores[
+ :, 0, : window_overlap - 1, 1 - window_overlap :
+ ]
+
+ # separate batch_size and num_heads dimensions again
+ diagonal_attention_scores = diagonal_attention_scores.view(
+ batch_size, num_heads, seq_len, 2 * window_overlap + 1
+ ).transpose(2, 1)
+
+ self._mask_invalid_locations(diagonal_attention_scores, window_overlap)
+ return diagonal_attention_scores
+
+ def _sliding_chunks_matmul_attn_probs_value(
+ self, attn_probs: torch.Tensor, value: torch.Tensor, window_overlap: int
+ ):
+ """
+ Same as _sliding_chunks_query_key_matmul but for attn_probs and value tensors. Returned tensor will be of the
+ same shape as `attn_probs`
+ """
+ batch_size, seq_len, num_heads, head_dim = value.size()
+
+ assert seq_len % (window_overlap * 2) == 0
+ assert attn_probs.size()[:3] == value.size()[:3]
+ assert attn_probs.size(3) == 2 * window_overlap + 1
+ chunks_count = torch.div(seq_len, window_overlap, rounding_mode="trunc") - 1
+ # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size 2 window overlap
+
+ chunked_attn_probs = attn_probs.transpose(1, 2).reshape(
+ batch_size * num_heads,
+ torch.div(seq_len, window_overlap, rounding_mode="trunc"),
+ window_overlap,
+ 2 * window_overlap + 1,
+ )
+
+ # group batch_size and num_heads dimensions into one
+ value = value.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+
+ # pad seq_len with w at the beginning of the sequence and another window overlap at the end
+ padded_value = nn.functional.pad(value, (0, 0, window_overlap, window_overlap), value=-1)
+
+ # chunk padded_value into chunks of size 3 window overlap and an overlap of size window overlap
+ chunked_value_size = (batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim)
+ chunked_value_stride = padded_value.stride()
+ chunked_value_stride = (
+ chunked_value_stride[0],
+ window_overlap * chunked_value_stride[1],
+ chunked_value_stride[1],
+ chunked_value_stride[2],
+ )
+ chunked_value = padded_value.as_strided(size=chunked_value_size, stride=chunked_value_stride)
+
+ chunked_attn_probs = self._pad_and_diagonalize(chunked_attn_probs)
+
+ context = torch.einsum("bcwd,bcdh->bcwh", (chunked_attn_probs, chunked_value))
+ return context.view(batch_size, num_heads, seq_len, head_dim).transpose(1, 2)
+
+ @staticmethod
+ def _get_global_attn_indices(is_index_global_attn):
+ """compute global attn indices required throughout forward pass"""
+ # helper variable
+ num_global_attn_indices = is_index_global_attn.long().sum(dim=1)
+
+ # max number of global attn indices in batch
+ max_num_global_attn_indices = num_global_attn_indices.max()
+
+ # indices of global attn
+ is_index_global_attn_nonzero = is_index_global_attn.nonzero(as_tuple=True)
+
+ # helper variable
+ is_local_index_global_attn = torch.arange(
+ max_num_global_attn_indices, device=is_index_global_attn.device
+ ) < num_global_attn_indices.unsqueeze(dim=-1)
+
+ # location of the non-padding values within global attention indices
+ is_local_index_global_attn_nonzero = is_local_index_global_attn.nonzero(as_tuple=True)
+
+ # location of the padding values within global attention indices
+ is_local_index_no_global_attn_nonzero = (is_local_index_global_attn == 0).nonzero(as_tuple=True)
+ return (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ )
+
+ def _concat_with_global_key_attn_probs(
+ self,
+ key_vectors,
+ query_vectors,
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ ):
+ batch_size = key_vectors.shape[0]
+
+ # create only global key vectors
+ key_vectors_only_global = key_vectors.new_zeros(
+ batch_size, max_num_global_attn_indices, self.num_heads, self.head_dim
+ )
+
+ key_vectors_only_global[is_local_index_global_attn_nonzero] = key_vectors[is_index_global_attn_nonzero]
+
+ # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
+ attn_probs_from_global_key = torch.einsum("blhd,bshd->blhs", (query_vectors, key_vectors_only_global))
+
+ # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets
+ attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3)
+ attn_probs_from_global_key[
+ is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, :
+ ] = torch.finfo(attn_probs_from_global_key.dtype).min
+ attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3)
+
+ return attn_probs_from_global_key
+
+ def _compute_attn_output_with_global_indices(
+ self,
+ value_vectors,
+ attn_probs,
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ ):
+ batch_size = attn_probs.shape[0]
+
+ # cut local attn probs to global only
+ attn_probs_only_global = attn_probs.narrow(-1, 0, max_num_global_attn_indices)
+ # get value vectors for global only
+ value_vectors_only_global = value_vectors.new_zeros(
+ batch_size, max_num_global_attn_indices, self.num_heads, self.head_dim
+ )
+ value_vectors_only_global[is_local_index_global_attn_nonzero] = value_vectors[is_index_global_attn_nonzero]
+
+ # use `matmul` because `einsum` crashes sometimes with fp16
+ # attn = torch.einsum('blhs,bshd->blhd', (selected_attn_probs, selected_v))
+ # compute attn output only global
+ attn_output_only_global = torch.matmul(
+ attn_probs_only_global.transpose(1, 2).clone(), value_vectors_only_global.transpose(1, 2).clone()
+ ).transpose(1, 2)
+
+ # reshape attn probs
+ attn_probs_without_global = attn_probs.narrow(
+ -1, max_num_global_attn_indices, attn_probs.size(-1) - max_num_global_attn_indices
+ ).contiguous()
+
+ # compute attn output with global
+ attn_output_without_global = self._sliding_chunks_matmul_attn_probs_value(
+ attn_probs_without_global, value_vectors, self.one_sided_attn_window_size
+ )
+ return attn_output_only_global + attn_output_without_global
+
+ def _compute_global_attn_output_from_hidden(
+ self,
+ hidden_states,
+ max_num_global_attn_indices,
+ layer_head_mask,
+ is_local_index_global_attn_nonzero,
+ is_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ is_index_masked,
+ ):
+ seq_len, batch_size = hidden_states.shape[:2]
+
+ # prepare global hidden states
+ global_attn_hidden_states = hidden_states.new_zeros(max_num_global_attn_indices, batch_size, self.embed_dim)
+ global_attn_hidden_states[is_local_index_global_attn_nonzero[::-1]] = hidden_states[
+ is_index_global_attn_nonzero[::-1]
+ ]
+
+ # global key, query, value
+ global_query_vectors_only_global = self.query_global(global_attn_hidden_states)
+ global_key_vectors = self.key_global(hidden_states)
+ global_value_vectors = self.value_global(hidden_states)
+
+ # normalize
+ global_query_vectors_only_global /= math.sqrt(self.head_dim)
+
+ # reshape
+ global_query_vectors_only_global = (
+ global_query_vectors_only_global.contiguous()
+ .view(max_num_global_attn_indices, batch_size * self.num_heads, self.head_dim)
+ .transpose(0, 1)
+ ) # (batch_size * self.num_heads, max_num_global_attn_indices, head_dim)
+ global_key_vectors = (
+ global_key_vectors.contiguous().view(-1, batch_size * self.num_heads, self.head_dim).transpose(0, 1)
+ ) # batch_size * self.num_heads, seq_len, head_dim)
+ global_value_vectors = (
+ global_value_vectors.contiguous().view(-1, batch_size * self.num_heads, self.head_dim).transpose(0, 1)
+ ) # batch_size * self.num_heads, seq_len, head_dim)
+
+ # compute attn scores
+ global_attn_scores = torch.bmm(global_query_vectors_only_global, global_key_vectors.transpose(1, 2))
+
+ assert list(global_attn_scores.size()) == [
+ batch_size * self.num_heads,
+ max_num_global_attn_indices,
+ seq_len,
+ ], (
+ "global_attn_scores have the wrong size. Size should be"
+ f" {(batch_size * self.num_heads, max_num_global_attn_indices, seq_len)}, but is"
+ f" {global_attn_scores.size()}."
+ )
+
+ global_attn_scores = global_attn_scores.view(batch_size, self.num_heads, max_num_global_attn_indices, seq_len)
+
+ # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets
+ global_attn_scores = global_attn_scores.transpose(1, 2)
+ global_attn_scores[
+ is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, :
+ ] = torch.finfo(global_attn_scores.dtype).min
+ global_attn_scores = global_attn_scores.transpose(1, 2)
+
+ global_attn_scores = global_attn_scores.masked_fill(
+ is_index_masked[:, None, None, :],
+ torch.finfo(global_attn_scores.dtype).min,
+ )
+
+ global_attn_scores = global_attn_scores.view(batch_size * self.num_heads, max_num_global_attn_indices, seq_len)
+
+ # compute global attn probs
+ global_attn_probs_float = nn.functional.softmax(
+ global_attn_scores, dim=-1, dtype=torch.float32
+ ) # use fp32 for numerical stability
+
+ # apply layer head masking
+ 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()}"
+ )
+ global_attn_probs_float = layer_head_mask.view(1, -1, 1, 1) * global_attn_probs_float.view(
+ batch_size, self.num_heads, max_num_global_attn_indices, seq_len
+ )
+ global_attn_probs_float = global_attn_probs_float.view(
+ batch_size * self.num_heads, max_num_global_attn_indices, seq_len
+ )
+
+ global_attn_probs = nn.functional.dropout(
+ global_attn_probs_float.type_as(global_attn_scores), p=self.dropout, training=self.training
+ )
+
+ # global attn output
+ global_attn_output = torch.bmm(global_attn_probs, global_value_vectors)
+
+ assert list(global_attn_output.size()) == [
+ batch_size * self.num_heads,
+ max_num_global_attn_indices,
+ self.head_dim,
+ ], (
+ "global_attn_output tensor has the wrong size. Size should be"
+ f" {(batch_size * self.num_heads, max_num_global_attn_indices, self.head_dim)}, but is"
+ f" {global_attn_output.size()}."
+ )
+
+ global_attn_probs = global_attn_probs.view(batch_size, self.num_heads, max_num_global_attn_indices, seq_len)
+ global_attn_output = global_attn_output.view(
+ batch_size, self.num_heads, max_num_global_attn_indices, self.head_dim
+ )
+ return global_attn_output, global_attn_probs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class LongformerSelfOutput(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 LongformerAttention(nn.Module):
+ def __init__(self, config, layer_id=0):
+ super().__init__()
+ self.self = LongformerSelfAttention(config, layer_id)
+ self.output = LongformerSelfOutput(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,
+ layer_head_mask=None,
+ is_index_masked=None,
+ is_index_global_attn=None,
+ is_global_attn=None,
+ output_attentions=False,
+ ):
+ self_outputs = self.self(
+ hidden_states,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ is_index_masked=is_index_masked,
+ is_index_global_attn=is_index_global_attn,
+ is_global_attn=is_global_attn,
+ output_attentions=output_attentions,
+ )
+ attn_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attn_output,) + self_outputs[1:]
+ return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate
+class LongformerIntermediate(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 LongformerOutput(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 LongformerLayer(nn.Module):
+ def __init__(self, config, layer_id=0):
+ super().__init__()
+ self.attention = LongformerAttention(config, layer_id)
+ self.intermediate = LongformerIntermediate(config)
+ self.output = LongformerOutput(config)
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ layer_head_mask=None,
+ is_index_masked=None,
+ is_index_global_attn=None,
+ is_global_attn=None,
+ output_attentions=False,
+ ):
+ self_attn_outputs = self.attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ is_index_masked=is_index_masked,
+ is_index_global_attn=is_index_global_attn,
+ is_global_attn=is_global_attn,
+ output_attentions=output_attentions,
+ )
+ attn_output = self_attn_outputs[0]
+ outputs = self_attn_outputs[1:]
+
+ layer_output = apply_chunking_to_forward(
+ self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attn_output
+ )
+ outputs = (layer_output,) + outputs
+ return outputs
+
+ def ff_chunk(self, attn_output):
+ intermediate_output = self.intermediate(attn_output)
+ layer_output = self.output(intermediate_output, attn_output)
+ return layer_output
+
+
+class LongformerEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([LongformerLayer(config, layer_id=i) for i in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ padding_len=0,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ ):
+ is_index_masked = attention_mask < 0
+ is_index_global_attn = attention_mask > 0
+
+ # Record `is_global_attn == True` to enable ONNX export
+ is_global_attn = is_index_global_attn.flatten().any().item()
+
+ all_hidden_states = () if output_hidden_states else None
+ all_attentions = () if output_attentions else None # All local attentions.
+ all_global_attentions = () if (output_attentions and is_global_attn) 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.layer)), (
+ f"The head_mask should be specified for {len(self.layer)} layers, but it is for {head_mask.size()[0]}."
+ )
+ for idx, 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,
+ attention_mask,
+ head_mask[idx] if head_mask is not None else None,
+ is_index_masked,
+ is_index_global_attn,
+ is_global_attn,
+ output_attentions,
+ )
+ else:
+ layer_outputs = layer_module(
+ hidden_states,
+ attention_mask=attention_mask,
+ layer_head_mask=head_mask[idx] if head_mask is not None else None,
+ is_index_masked=is_index_masked,
+ is_index_global_attn=is_index_global_attn,
+ is_global_attn=is_global_attn,
+ output_attentions=output_attentions,
+ )
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ # bzs x seq_len x num_attn_heads x (num_global_attn + attention_window_len + 1) => bzs x num_attn_heads x seq_len x (num_global_attn + attention_window_len + 1)
+ all_attentions = all_attentions + (layer_outputs[1].transpose(1, 2),)
+
+ if is_global_attn:
+ # bzs x num_attn_heads x num_global_attn x seq_len => bzs x num_attn_heads x seq_len x num_global_attn
+ all_global_attentions = all_global_attentions + (layer_outputs[2].transpose(2, 3),)
+
+ # Add last layer
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ # undo padding if necessary
+ # unpad `hidden_states` because the calling function is expecting a length == input_ids.size(1)
+ hidden_states = hidden_states[:, : hidden_states.shape[1] - padding_len]
+ if output_hidden_states:
+ all_hidden_states = tuple([state[:, : state.shape[1] - padding_len] for state in all_hidden_states])
+
+ if output_attentions:
+ all_attentions = tuple([state[:, :, : state.shape[2] - padding_len, :] for state in all_attentions])
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, all_hidden_states, all_attentions, all_global_attentions] if v is not None
+ )
+ return LongformerBaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_attentions,
+ global_attentions=all_global_attentions,
+ )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler
+class LongformerPooler(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.roberta.modeling_roberta.RobertaLMHead with Roberta->Longformer
+class LongformerLMHead(nn.Module):
+ """Longformer 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
+
+
+class LongformerPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = LongformerConfig
+ base_model_prefix = "longformer"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["LongformerSelfAttention"]
+
+ 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)
+
+
+LONGFORMER_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 ([`LongformerConfig`]): 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.
+"""
+
+LONGFORMER_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)
+ global_attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
+ Mask to decide the attention given on each token, local attention or global attention. Tokens with global
+ attention attends to all other tokens, and all other tokens attend to them. This is important for
+ task-specific finetuning because it makes the model more flexible at representing the task. For example,
+ for classification, the token should be given global attention. For QA, all question tokens should also
+ have global attention. Please refer to the [Longformer paper](https://arxiv.org/abs/2004.05150) for more
+ details. Mask values selected in `[0, 1]`:
+
+ - 0 for local attention (a sliding window attention),
+ - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them).
+
+ head_mask (`torch.Tensor` of shape `(num_layers, num_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 `(num_layers, num_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**.
+
+ 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 Longformer Model outputting raw hidden-states without any specific head on top.",
+ LONGFORMER_START_DOCSTRING,
+)
+class LongformerModel(LongformerPreTrainedModel):
+ """
+ This class copied code from [`RobertaModel`] and overwrote standard self-attention with longformer self-attention
+ to provide the ability to process long sequences following the self-attention approach described in [Longformer:
+ the Long-Document Transformer](https://arxiv.org/abs/2004.05150) by Iz Beltagy, Matthew E. Peters, and Arman Cohan.
+ Longformer self-attention combines a local (sliding window) and global attention to extend to long documents
+ without the O(n^2) increase in memory and compute.
+
+ The self-attention module `LongformerSelfAttention` implemented here supports the combination of local and global
+ attention but it lacks support for autoregressive attention and dilated attention. Autoregressive and dilated
+ attention are more relevant for autoregressive language modeling than finetuning on downstream tasks. Future
+ release will add support for autoregressive attention, but the support for dilated attention requires a custom CUDA
+ kernel to be memory and compute efficient.
+
+ """
+
+ def __init__(self, config, add_pooling_layer=True):
+ super().__init__(config)
+ self.config = config
+
+ if isinstance(config.attention_window, int):
+ assert config.attention_window % 2 == 0, "`config.attention_window` has to be an even value"
+ assert config.attention_window > 0, "`config.attention_window` has to be positive"
+ config.attention_window = [config.attention_window] * config.num_hidden_layers # one value per layer
+ else:
+ assert len(config.attention_window) == config.num_hidden_layers, (
+ "`len(config.attention_window)` should equal `config.num_hidden_layers`. "
+ f"Expected {config.num_hidden_layers}, given {len(config.attention_window)}"
+ )
+
+ self.embeddings = LongformerEmbeddings(config)
+ self.encoder = LongformerEncoder(config)
+ self.pooler = LongformerPooler(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)
+
+ def _pad_to_window_size(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ token_type_ids: torch.Tensor,
+ position_ids: torch.Tensor,
+ inputs_embeds: torch.Tensor,
+ pad_token_id: int,
+ ):
+ """A helper function to pad tokens and mask to work with implementation of Longformer self-attention."""
+ # padding
+ attention_window = (
+ self.config.attention_window
+ if isinstance(self.config.attention_window, int)
+ else max(self.config.attention_window)
+ )
+
+ assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"
+ input_shape = input_ids.shape if input_ids is not None else inputs_embeds.shape
+ batch_size, seq_len = input_shape[:2]
+
+ padding_len = (attention_window - seq_len % attention_window) % attention_window
+
+ # this path should be recorded in the ONNX export, it is fine with padding_len == 0 as well
+ if padding_len > 0:
+ logger.warning_once(
+ f"Input ids are automatically padded to be a multiple of `config.attention_window`: {attention_window}"
+ )
+ if input_ids is not None:
+ input_ids = nn.functional.pad(input_ids, (0, padding_len), value=pad_token_id)
+ if position_ids is not None:
+ # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
+ position_ids = nn.functional.pad(position_ids, (0, padding_len), value=pad_token_id)
+ if inputs_embeds is not None:
+ input_ids_padding = inputs_embeds.new_full(
+ (batch_size, padding_len),
+ self.config.pad_token_id,
+ dtype=torch.long,
+ )
+ inputs_embeds_padding = self.embeddings(input_ids_padding)
+ inputs_embeds = torch.cat([inputs_embeds, inputs_embeds_padding], dim=-2)
+
+ attention_mask = nn.functional.pad(
+ attention_mask, (0, padding_len), value=0
+ ) # no attention on the padding tokens
+ token_type_ids = nn.functional.pad(token_type_ids, (0, padding_len), value=0) # pad with token_type_id = 0
+
+ return padding_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds
+
+ def _merge_to_attention_mask(self, attention_mask: torch.Tensor, global_attention_mask: torch.Tensor):
+ # longformer self attention expects attention mask to have 0 (no attn), 1 (local attn), 2 (global attn)
+ # (global_attention_mask + 1) => 1 for local attention, 2 for global attention
+ # => final attention_mask => 0 for no attention, 1 for local attention 2 for global attention
+ if attention_mask is not None:
+ attention_mask = attention_mask * (global_attention_mask + 1)
+ else:
+ # simply use `global_attention_mask` as `attention_mask`
+ # if no `attention_mask` is given
+ attention_mask = global_attention_mask + 1
+ return attention_mask
+
+ @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=LongformerBaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_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, LongformerBaseModelOutputWithPooling]:
+ r"""
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import LongformerModel, AutoTokenizer
+
+ >>> model = LongformerModel.from_pretrained("allenai/longformer-base-4096")
+ >>> tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
+
+ >>> SAMPLE_TEXT = " ".join(["Hello world! "] * 1000) # long input document
+ >>> input_ids = torch.tensor(tokenizer.encode(SAMPLE_TEXT)).unsqueeze(0) # batch of size 1
+
+ >>> attention_mask = torch.ones(
+ ... input_ids.shape, dtype=torch.long, device=input_ids.device
+ ... ) # initialize to local attention
+ >>> global_attention_mask = torch.zeros(
+ ... input_ids.shape, dtype=torch.long, device=input_ids.device
+ ... ) # initialize to global attention to be deactivated for all tokens
+ >>> global_attention_mask[
+ ... :,
+ ... [
+ ... 1,
+ ... 4,
+ ... 21,
+ ... ],
+ ... ] = 1 # Set global attention to random tokens for the sake of this example
+ >>> # Usually, set global attention based on the task. For example,
+ >>> # classification: the token
+ >>> # QA: question tokens
+ >>> # LM: potentially on the beginning of sentences and paragraphs
+ >>> outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
+ >>> sequence_output = outputs.last_hidden_state
+ >>> pooled_output = outputs.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 token_type_ids is None:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+ # merge `global_attention_mask` and `attention_mask`
+ if global_attention_mask is not None:
+ attention_mask = self._merge_to_attention_mask(attention_mask, global_attention_mask)
+
+ padding_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds = self._pad_to_window_size(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ pad_token_id=self.config.pad_token_id,
+ )
+
+ # 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)[
+ :, 0, 0, :
+ ]
+
+ 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,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ padding_len=padding_len,
+ 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 LongformerBaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ global_attentions=encoder_outputs.global_attentions,
+ )
+
+
+@add_start_docstrings("""Longformer Model with a `language modeling` head on top.""", LONGFORMER_START_DOCSTRING)
+class LongformerForMaskedLM(LongformerPreTrainedModel):
+ _tied_weights_keys = ["lm_head.decoder"]
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.longformer = LongformerModel(config, add_pooling_layer=False)
+ self.lm_head = LongformerLMHead(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(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=LongformerMaskedLMOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_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, LongformerMaskedLMOutput]:
+ 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.
+
+ Returns:
+
+ Mask filling example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LongformerForMaskedLM
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
+ >>> model = LongformerForMaskedLM.from_pretrained("allenai/longformer-base-4096")
+ ```
+
+ Let's try a very long input.
+
+ ```python
+ >>> TXT = (
+ ... "My friends are but they eat too many carbs."
+ ... + " That's why I decide not to eat with them." * 300
+ ... )
+ >>> input_ids = tokenizer([TXT], return_tensors="pt")["input_ids"]
+ >>> logits = model(input_ids).logits
+
+ >>> masked_index = (input_ids[0] == tokenizer.mask_token_id).nonzero().item()
+ >>> probs = logits[0, masked_index].softmax(dim=0)
+ >>> values, predictions = probs.topk(5)
+
+ >>> tokenizer.decode(predictions).split()
+ ['healthy', 'skinny', 'thin', 'good', 'vegetarian']
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.longformer(
+ input_ids,
+ attention_mask=attention_mask,
+ global_attention_mask=global_attention_mask,
+ head_mask=head_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.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 LongformerMaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=prediction_scores,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ Longformer Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+ pooled output) e.g. for GLUE tasks.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class LongformerForSequenceClassification(LongformerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.config = config
+
+ self.longformer = LongformerModel(config, add_pooling_layer=False)
+ self.classifier = LongformerClassificationHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint="jpwahle/longformer-base-plagiarism-detection",
+ output_type=LongformerSequenceClassifierOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output="'ORIGINAL'",
+ expected_loss=5.44,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_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, LongformerSequenceClassifierOutput]:
+ 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
+
+ if global_attention_mask is None:
+ logger.warning_once("Initializing global attention on CLS token...")
+ global_attention_mask = torch.zeros_like(input_ids)
+ # global attention on cls token
+ global_attention_mask[:, 0] = 1
+
+ outputs = self.longformer(
+ input_ids,
+ attention_mask=attention_mask,
+ global_attention_mask=global_attention_mask,
+ head_mask=head_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.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 LongformerSequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+
+class LongformerClassificationHead(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, hidden_states, **kwargs):
+ hidden_states = hidden_states[:, 0, :] # take token (equiv. to [CLS])
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.dense(hidden_states)
+ hidden_states = torch.tanh(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ output = self.out_proj(hidden_states)
+ return output
+
+
+@add_start_docstrings(
+ """
+ Longformer Model with a span classification head on top for extractive question-answering tasks like SQuAD /
+ TriviaQA (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class LongformerForQuestionAnswering(LongformerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.longformer = LongformerModel(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(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=LongformerQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_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, LongformerQuestionAnsweringModelOutput]:
+ 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.
+
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, LongformerForQuestionAnswering
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-large-4096-finetuned-triviaqa")
+ >>> model = LongformerForQuestionAnswering.from_pretrained("allenai/longformer-large-4096-finetuned-triviaqa")
+
+ >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+ >>> encoding = tokenizer(question, text, return_tensors="pt")
+ >>> input_ids = encoding["input_ids"]
+
+ >>> # default is local attention everywhere
+ >>> # the forward method will automatically set global attention on question tokens
+ >>> attention_mask = encoding["attention_mask"]
+
+ >>> outputs = model(input_ids, attention_mask=attention_mask)
+ >>> start_logits = outputs.start_logits
+ >>> end_logits = outputs.end_logits
+ >>> all_tokens = tokenizer.convert_ids_to_tokens(input_ids[0].tolist())
+
+ >>> answer_tokens = all_tokens[torch.argmax(start_logits) : torch.argmax(end_logits) + 1]
+ >>> answer = tokenizer.decode(
+ ... tokenizer.convert_tokens_to_ids(answer_tokens)
+ ... ) # remove space prepending space token
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if global_attention_mask is None:
+ if input_ids is None:
+ logger.warning(
+ "It is not possible to automatically generate the `global_attention_mask` because input_ids is"
+ " None. Please make sure that it is correctly set."
+ )
+ else:
+ # set global attention on question tokens automatically
+ global_attention_mask = _compute_global_attention_mask(input_ids, self.config.sep_token_id)
+
+ outputs = self.longformer(
+ input_ids,
+ attention_mask=attention_mask,
+ global_attention_mask=global_attention_mask,
+ head_mask=head_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[2:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return LongformerQuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ Longformer 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.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class LongformerForTokenClassification(LongformerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.longformer = LongformerModel(config, add_pooling_layer=False)
+ 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(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint="brad1141/Longformer-finetuned-norm",
+ output_type=LongformerTokenClassifierOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output=(
+ "['Evidence', 'Evidence', 'Evidence', 'Evidence', 'Evidence', 'Evidence', 'Evidence', 'Evidence',"
+ " 'Evidence', 'Evidence', 'Evidence', 'Evidence']"
+ ),
+ expected_loss=0.63,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_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, LongformerTokenClassifierOutput]:
+ 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.longformer(
+ input_ids,
+ attention_mask=attention_mask,
+ global_attention_mask=global_attention_mask,
+ head_mask=head_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()
+
+ 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 LongformerTokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ Longformer 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.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class LongformerForMultipleChoice(LongformerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.longformer = LongformerModel(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(
+ LONGFORMER_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
+ )
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=LongformerMultipleChoiceModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ global_attention_mask: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ labels: 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, LongformerMultipleChoiceModelOutput]:
+ 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)
+ """
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # set global attention on question tokens
+ if global_attention_mask is None and input_ids is not None:
+ logger.warning_once("Initializing global attention on multiple choice...")
+ # put global attention on all tokens after `config.sep_token_id`
+ global_attention_mask = torch.stack(
+ [
+ _compute_global_attention_mask(input_ids[:, i], self.config.sep_token_id, before_sep_token=False)
+ for i in range(num_choices)
+ ],
+ dim=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_global_attention_mask = (
+ global_attention_mask.view(-1, global_attention_mask.size(-1))
+ if global_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.longformer(
+ flat_input_ids,
+ position_ids=flat_position_ids,
+ token_type_ids=flat_token_type_ids,
+ attention_mask=flat_attention_mask,
+ global_attention_mask=flat_global_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 LongformerMultipleChoiceModelOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+
+__all__ = [
+ "LongformerForMaskedLM",
+ "LongformerForMultipleChoice",
+ "LongformerForQuestionAnswering",
+ "LongformerForSequenceClassification",
+ "LongformerForTokenClassification",
+ "LongformerModel",
+ "LongformerPreTrainedModel",
+ "LongformerSelfAttention",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/modeling_tf_longformer.py b/lib/python3.12/site-packages/transformers/models/longformer/modeling_tf_longformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..9280838de0b9da8b3880050825a178b14232e1fd
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/modeling_tf_longformer.py
@@ -0,0 +1,2786 @@
+# coding=utf-8
+# Copyright 2020 The Allen Institute for AI team 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.
+"""Tensorflow Longformer 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_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,
+)
+from .configuration_longformer import LongformerConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "allenai/longformer-base-4096"
+_CONFIG_FOR_DOC = "LongformerConfig"
+
+LARGE_NEGATIVE = -1e8
+
+
+@dataclass
+class TFLongformerBaseModelOutput(ModelOutput):
+ """
+ Base class for Longformer's outputs, with potential hidden states, local and global attentions.
+
+ 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.
+ 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ last_hidden_state: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerBaseModelOutputWithPooling(ModelOutput):
+ """
+ Base class for Longformer's outputs that also contains a pooling of the last hidden states.
+
+ 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)`):
+ Last layer hidden-state of the first token of the sequence (classification token) further processed by 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(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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ last_hidden_state: Optional[tf.Tensor] = None
+ pooler_output: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerMaskedLMOutput(ModelOutput):
+ """
+ Base class for masked language models outputs.
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Masked language modeling (MLM) loss.
+ 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerQuestionAnsweringModelOutput(ModelOutput):
+ """
+ Base class for outputs of question answering Longformer models.
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+ start_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
+ Span-start scores (before SoftMax).
+ end_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
+ Span-end scores (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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: tf.Tensor | None = None
+ start_logits: Optional[tf.Tensor] = None
+ end_logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerSequenceClassifierOutput(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).
+ 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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerMultipleChoiceModelOutput(ModelOutput):
+ """
+ Base class for outputs of multiple choice models.
+
+ Args:
+ loss (`tf.Tensor` of shape *(1,)*, *optional*, returned when `labels` is provided):
+ Classification loss.
+ logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
+ *num_choices* is the second dimension of the input tensors. (see *input_ids* above).
+
+ Classification scores (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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFLongformerTokenClassifierOutput(ModelOutput):
+ """
+ Base class for outputs of token classification models.
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :
+ Classification loss.
+ logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.num_labels)`):
+ Classification scores (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, x +
+ attention_window + 1)`, where `x` is the number of tokens with global attention mask.
+
+ Local attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token in the sequence to every token with
+ global attention (first `x` values) and to every token in the attention window (remaining `attention_window
+ + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the
+ remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a
+ token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding
+ (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.
+ If the attention window contains a token with global attention, the attention weight at the corresponding
+ index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global
+ attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be
+ accessed from `global_attentions`.
+ global_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, x)`, where `x`
+ is the number of tokens with global attention mask.
+
+ Global attentions weights after the attention softmax, used to compute the weighted average in the
+ self-attention heads. Those are the attention weights from every token with global attention to every token
+ in the sequence.
+ """
+
+ loss: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ global_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):
+ """
+ Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is
+ True` else after `sep_token_id`.
+ """
+ assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"
+ question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]
+ # bool attention mask with True in locations of global attention
+ attention_mask = tf.expand_dims(tf.range(input_ids_shape[1], dtype=tf.int64), axis=0)
+ attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))
+ if before_sep_token is True:
+ question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))
+ attention_mask = tf.cast(attention_mask < question_end_index, dtype=question_end_index.dtype)
+ else:
+ # last token is separation token and should not be counted and in the middle are two separation tokens
+ question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))
+ attention_mask = tf.cast(
+ attention_mask > question_end_index,
+ dtype=question_end_index.dtype,
+ ) * tf.cast(attention_mask < input_ids_shape[-1], dtype=question_end_index.dtype)
+
+ return attention_mask
+
+
+# Copied from transformers.models.roberta.modeling_tf_roberta.TFRobertaLMHead with Roberta->Longformer
+class TFLongformerLMHead(keras.layers.Layer):
+ """Longformer Head for masked language modeling."""
+
+ def __init__(self, config, input_embeddings, **kwargs):
+ super().__init__(**kwargs)
+
+ self.config = config
+ self.hidden_size = config.hidden_size
+ 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")
+ self.act = get_tf_activation("gelu")
+
+ # The output weights are the same as the input embeddings, but there is
+ # an output-only bias for each token.
+ self.decoder = 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, "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])
+
+ def get_output_embeddings(self):
+ return self.decoder
+
+ def set_output_embeddings(self, value):
+ self.decoder.weight = value
+ self.decoder.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):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.layer_norm(hidden_states)
+
+ # project back to size of vocabulary with bias
+ 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.decoder.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 TFLongformerEmbeddings(keras.layers.Layer):
+ """
+ Same as BertEmbeddings with a tiny tweak for positional embeddings indexing and some extra casting.
+ """
+
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+
+ self.padding_idx = 1
+ self.config = config
+ self.hidden_size = config.hidden_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.hidden_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.hidden_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.hidden_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.hidden_size])
+
+ def create_position_ids_from_input_ids(self, input_ids, 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:
+ input_ids: tf.Tensor
+ Returns: tf.Tensor
+ """
+ mask = tf.cast(tf.math.not_equal(input_ids, self.padding_idx), dtype=input_ids.dtype)
+ incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask
+
+ return incremental_indices + self.padding_idx
+
+ def call(
+ self,
+ input_ids=None,
+ position_ids=None,
+ token_type_ids=None,
+ inputs_embeds=None,
+ past_key_values_length=0,
+ 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 token_type_ids is None:
+ token_type_ids = tf.cast(tf.fill(dims=input_shape, value=0), tf.int64)
+
+ 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 = self.create_position_ids_from_input_ids(
+ input_ids=input_ids, past_key_values_length=past_key_values_length
+ )
+ else:
+ position_ids = tf.expand_dims(
+ tf.range(start=self.padding_idx + 1, limit=input_shape[-1] + self.padding_idx + 1, dtype=tf.int64),
+ 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
+
+
+# Copied from transformers.models.bert.modeling_tf_bert.TFBertIntermediate with Bert->Longformer
+class TFLongformerIntermediate(keras.layers.Layer):
+ def __init__(self, config: LongformerConfig, **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->Longformer
+class TFLongformerOutput(keras.layers.Layer):
+ def __init__(self, config: LongformerConfig, **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.TFBertPooler with Bert->Longformer
+class TFLongformerPooler(keras.layers.Layer):
+ def __init__(self, config: LongformerConfig, **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.bert.modeling_tf_bert.TFBertSelfOutput with Bert->Longformer
+class TFLongformerSelfOutput(keras.layers.Layer):
+ def __init__(self, config: LongformerConfig, **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])
+
+
+class TFLongformerSelfAttention(keras.layers.Layer):
+ def __init__(self, config, layer_id, **kwargs):
+ super().__init__(**kwargs)
+ self.config = config
+
+ 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_heads = config.num_attention_heads
+ self.head_dim = int(config.hidden_size / config.num_attention_heads)
+ self.embed_dim = config.hidden_size
+ self.query = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="query",
+ )
+ self.key = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="key",
+ )
+ self.value = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="value",
+ )
+
+ # separate projection layers for tokens with global attention
+ self.query_global = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="query_global",
+ )
+ self.key_global = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="key_global",
+ )
+ self.value_global = keras.layers.Dense(
+ self.embed_dim,
+ kernel_initializer=get_initializer(config.initializer_range),
+ name="value_global",
+ )
+ self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
+ self.global_dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
+ self.layer_id = layer_id
+ attention_window = config.attention_window[self.layer_id]
+
+ assert attention_window % 2 == 0, (
+ f"`attention_window` for layer {self.layer_id} has to be an even value. Given {attention_window}"
+ )
+ assert attention_window > 0, (
+ f"`attention_window` for layer {self.layer_id} has to be positive. Given {attention_window}"
+ )
+
+ self.one_sided_attn_window_size = attention_window // 2
+
+ def build(self, input_shape=None):
+ if not self.built:
+ with tf.name_scope("query_global"):
+ self.query_global.build((self.config.hidden_size,))
+ with tf.name_scope("key_global"):
+ self.key_global.build((self.config.hidden_size,))
+ with tf.name_scope("value_global"):
+ self.value_global.build((self.config.hidden_size,))
+
+ 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, "query_global", None) is not None:
+ with tf.name_scope(self.query_global.name):
+ self.query_global.build([None, None, self.config.hidden_size])
+ if getattr(self, "key_global", None) is not None:
+ with tf.name_scope(self.key_global.name):
+ self.key_global.build([None, None, self.config.hidden_size])
+ if getattr(self, "value_global", None) is not None:
+ with tf.name_scope(self.value_global.name):
+ self.value_global.build([None, None, self.config.hidden_size])
+
+ def call(
+ self,
+ inputs,
+ training=False,
+ ):
+ """
+ LongformerSelfAttention expects *len(hidden_states)* to be multiple of *attention_window*. Padding to
+ *attention_window* happens in LongformerModel.forward to avoid redoing the padding on each layer.
+
+ The *attention_mask* is changed in [`LongformerModel.forward`] from 0, 1, 2 to:
+
+ - -10000: no attention
+ - 0: local attention
+ - +10000: global attention
+ """
+ # retrieve input args
+ (
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ is_index_masked,
+ is_index_global_attn,
+ is_global_attn,
+ ) = inputs
+
+ # project hidden states
+ query_vectors = self.query(hidden_states)
+ key_vectors = self.key(hidden_states)
+ value_vectors = self.value(hidden_states)
+ batch_size, seq_len, embed_dim = shape_list(hidden_states)
+
+ tf.debugging.assert_equal(
+ embed_dim,
+ self.embed_dim,
+ message=f"hidden_states should have embed_dim = {self.embed_dim}, but has {embed_dim}",
+ )
+
+ # normalize query
+ query_vectors /= tf.math.sqrt(tf.cast(self.head_dim, dtype=query_vectors.dtype))
+ query_vectors = tf.reshape(query_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))
+ key_vectors = tf.reshape(key_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))
+
+ # attn_probs = (batch_size, seq_len, num_heads, window*2+1)
+ attn_scores = self._sliding_chunks_query_key_matmul(
+ query_vectors, key_vectors, self.one_sided_attn_window_size
+ )
+
+ # values to pad for attention probs
+ remove_from_windowed_attention_mask = attention_mask != 0
+ # cast to fp32/fp16 then replace 1's with -inf
+ float_mask = tf.cast(remove_from_windowed_attention_mask, dtype=query_vectors.dtype) * LARGE_NEGATIVE
+
+ # diagonal mask with zeros everywhere and -inf inplace of padding
+ diagonal_mask = self._sliding_chunks_query_key_matmul(
+ tf.ones(shape_list(attention_mask)),
+ float_mask,
+ self.one_sided_attn_window_size,
+ )
+
+ # pad local attention probs
+ attn_scores += diagonal_mask
+
+ tf.debugging.assert_equal(
+ shape_list(attn_scores),
+ [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1],
+ message=(
+ f"attn_probs should be of size ({batch_size}, {seq_len}, {self.num_heads},"
+ f" {self.one_sided_attn_window_size * 2 + 1}), but is of size {shape_list(attn_scores)}"
+ ),
+ )
+
+ # compute global attn indices required through out forward fn
+ (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ ) = self._get_global_attn_indices(is_index_global_attn)
+
+ # this function is only relevant for global attention
+ if is_global_attn:
+ attn_scores = self._concat_with_global_key_attn_probs(
+ attn_scores=attn_scores,
+ query_vectors=query_vectors,
+ key_vectors=key_vectors,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ )
+
+ attn_probs = stable_softmax(attn_scores, axis=-1)
+
+ # softmax sometimes inserts NaN if all positions are masked, replace them with 0
+ # Make sure to create a mask with the proper shape:
+ # if is_global_attn==True => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1]
+ # if is_global_attn==False => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]
+ if is_global_attn:
+ masked_index = tf.tile(
+ is_index_masked[:, :, None, None],
+ (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1),
+ )
+ else:
+ masked_index = tf.tile(
+ is_index_masked[:, :, None, None],
+ (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + 1),
+ )
+ attn_probs = tf.where(
+ masked_index,
+ tf.zeros(shape_list(masked_index), dtype=attn_probs.dtype),
+ attn_probs,
+ )
+
+ if layer_head_mask is not None:
+ tf.debugging.assert_equal(
+ shape_list(layer_head_mask),
+ [self.num_heads],
+ message=(
+ f"Head mask for a single layer should be of size {(self.num_heads)}, but is"
+ f" {shape_list(layer_head_mask)}"
+ ),
+ )
+
+ attn_probs = tf.reshape(layer_head_mask, (1, 1, -1, 1)) * attn_probs
+
+ # apply dropout
+ attn_probs = self.dropout(attn_probs, training=training)
+ value_vectors = tf.reshape(value_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))
+
+ # if global attention, compute sum of global and local attn
+
+ if is_global_attn:
+ attn_output = self._compute_attn_output_with_global_indices(
+ value_vectors=value_vectors,
+ attn_probs=attn_probs,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ )
+ else:
+ attn_output = self._sliding_chunks_matmul_attn_probs_value(
+ attn_probs, value_vectors, self.one_sided_attn_window_size
+ )
+
+ tf.debugging.assert_equal(
+ shape_list(attn_output), [batch_size, seq_len, self.num_heads, self.head_dim], message="Unexpected size"
+ )
+
+ attn_output = tf.reshape(attn_output, (batch_size, seq_len, embed_dim))
+
+ # compute value for global attention and overwrite to attention output
+ if is_global_attn:
+ attn_output, global_attn_probs = self._compute_global_attn_output_from_hidden(
+ attn_output=attn_output,
+ hidden_states=hidden_states,
+ max_num_global_attn_indices=max_num_global_attn_indices,
+ layer_head_mask=layer_head_mask,
+ is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
+ is_index_global_attn_nonzero=is_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
+ is_index_masked=is_index_masked,
+ training=training,
+ )
+ else:
+ # Leave attn_output unchanged
+ global_attn_probs = tf.zeros((batch_size, self.num_heads, max_num_global_attn_indices, seq_len))
+
+ # make sure that local attention probabilities are set to 0 for indices of global attn
+ # Make sure to create a mask with the proper shape:
+ # if is_global_attn==True => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1]
+ # if is_global_attn==False => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]
+ if is_global_attn:
+ masked_global_attn_index = tf.tile(
+ is_index_global_attn[:, :, None, None],
+ (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1),
+ )
+ else:
+ masked_global_attn_index = tf.tile(
+ is_index_global_attn[:, :, None, None],
+ (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + 1),
+ )
+ attn_probs = tf.where(
+ masked_global_attn_index,
+ tf.zeros(shape_list(masked_global_attn_index), dtype=attn_probs.dtype),
+ attn_probs,
+ )
+
+ outputs = (attn_output, attn_probs, global_attn_probs)
+
+ return outputs
+
+ def _sliding_chunks_query_key_matmul(self, query, key, window_overlap):
+ """
+ Matrix multiplication of query and key tensors using with a sliding window attention pattern. This
+ implementation splits the input into overlapping chunks of size 2w (e.g. 512 for pretrained Longformer) with an
+ overlap of size window_overlap
+ """
+ batch_size, seq_len, num_heads, head_dim = shape_list(query)
+
+ tf.debugging.assert_equal(
+ seq_len % (window_overlap * 2),
+ 0,
+ message=f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}",
+ )
+ tf.debugging.assert_equal(
+ shape_list(query),
+ shape_list(key),
+ message=(
+ f"Shape of query and key should be equal, but got query: {shape_list(query)} and key:"
+ f" {shape_list(key)}"
+ ),
+ )
+
+ chunks_count = seq_len // window_overlap - 1
+
+ # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size window_overlap * 2
+ query = tf.reshape(
+ tf.transpose(query, (0, 2, 1, 3)),
+ (batch_size * num_heads, seq_len, head_dim),
+ )
+ key = tf.reshape(tf.transpose(key, (0, 2, 1, 3)), (batch_size * num_heads, seq_len, head_dim))
+ chunked_query = self._chunk(query, window_overlap)
+ chunked_key = self._chunk(key, window_overlap)
+
+ # matrix multiplication
+ # bcxd: batch_size * num_heads x chunks x 2window_overlap x head_dim
+ # bcyd: batch_size * num_heads x chunks x 2window_overlap x head_dim
+ # bcxy: batch_size * num_heads x chunks x 2window_overlap x 2window_overlap
+ chunked_query = tf.cast(chunked_query, dtype=chunked_key.dtype)
+ chunked_attention_scores = tf.einsum("bcxd,bcyd->bcxy", chunked_query, chunked_key) # multiply
+
+ # convert diagonals into columns
+ paddings = tf.convert_to_tensor([[0, 0], [0, 0], [0, 1], [0, 0]])
+ diagonal_chunked_attention_scores = self._pad_and_transpose_last_two_dims(chunked_attention_scores, paddings)
+
+ # allocate space for the overall attention matrix where the chunks are combined. The last dimension
+ # has (window_overlap * 2 + 1) columns. The first (window_overlap) columns are the window_overlap lower triangles (attention from a word to
+ # window_overlap previous words). The following column is attention score from each word to itself, then
+ # followed by window_overlap columns for the upper triangle.
+
+ # copy parts from diagonal_chunked_attention_scores into the combined matrix of attentions
+ # - copying the main diagonal and the upper triangle
+ # TODO: This code is most likely not very efficient and should be improved
+ diagonal_attn_scores_up_triang = tf.concat(
+ [
+ diagonal_chunked_attention_scores[:, :, :window_overlap, : window_overlap + 1],
+ diagonal_chunked_attention_scores[:, -1:, window_overlap:, : window_overlap + 1],
+ ],
+ axis=1,
+ )
+
+ # - copying the lower triangle
+ diagonal_attn_scores_low_triang = tf.concat(
+ [
+ tf.zeros(
+ (batch_size * num_heads, 1, window_overlap, window_overlap),
+ dtype=diagonal_chunked_attention_scores.dtype,
+ ),
+ diagonal_chunked_attention_scores[:, :, -(window_overlap + 1) : -1, window_overlap + 1 :],
+ ],
+ axis=1,
+ )
+ diagonal_attn_scores_first_chunk = tf.concat(
+ [
+ tf.roll(
+ diagonal_chunked_attention_scores,
+ shift=[1, window_overlap],
+ axis=[2, 3],
+ )[:, :, :window_overlap, :window_overlap],
+ tf.zeros(
+ (batch_size * num_heads, 1, window_overlap, window_overlap),
+ dtype=diagonal_chunked_attention_scores.dtype,
+ ),
+ ],
+ axis=1,
+ )
+ first_chunk_mask = (
+ tf.tile(
+ tf.range(chunks_count + 1, dtype=tf.int64)[None, :, None, None],
+ (batch_size * num_heads, 1, window_overlap, window_overlap),
+ )
+ < 1
+ )
+ diagonal_attn_scores_low_triang = tf.where(
+ first_chunk_mask,
+ diagonal_attn_scores_first_chunk,
+ diagonal_attn_scores_low_triang,
+ )
+
+ # merging upper and lower triangle
+ diagonal_attention_scores = tf.concat(
+ [diagonal_attn_scores_low_triang, diagonal_attn_scores_up_triang], axis=-1
+ )
+
+ # separate batch_size and num_heads dimensions again
+ diagonal_attention_scores = tf.transpose(
+ tf.reshape(
+ diagonal_attention_scores,
+ (batch_size, num_heads, seq_len, 2 * window_overlap + 1),
+ ),
+ (0, 2, 1, 3),
+ )
+
+ diagonal_attention_scores = self._mask_invalid_locations(diagonal_attention_scores, window_overlap)
+
+ return diagonal_attention_scores
+
+ @staticmethod
+ def _mask_invalid_locations(input_tensor, window_overlap):
+ # create correct upper triangle bool mask
+ mask_2d_upper = tf.reverse(
+ tf.linalg.band_part(tf.ones(shape=(window_overlap, window_overlap + 1)), -1, 0),
+ axis=[0],
+ )
+
+ # pad to full matrix
+ padding = tf.convert_to_tensor(
+ [[0, shape_list(input_tensor)[1] - window_overlap], [0, shape_list(input_tensor)[3] - window_overlap - 1]]
+ )
+
+ # create lower mask
+ mask_2d = tf.pad(mask_2d_upper, padding)
+
+ # combine with upper mask
+ mask_2d = mask_2d + tf.reverse(mask_2d, axis=[0, 1])
+
+ # broadcast to full matrix
+ mask_4d = tf.tile(mask_2d[None, :, None, :], (shape_list(input_tensor)[0], 1, 1, 1))
+
+ # inf tensor used for masking
+ inf_tensor = -float("inf") * tf.ones_like(input_tensor)
+
+ # mask
+ input_tensor = tf.where(tf.math.greater(mask_4d, 0), inf_tensor, input_tensor)
+
+ return input_tensor
+
+ def _sliding_chunks_matmul_attn_probs_value(self, attn_probs, value, window_overlap):
+ """
+ Same as _sliding_chunks_query_key_matmul but for attn_probs and value tensors. Returned tensor will be of the
+ same shape as `attn_probs`
+ """
+
+ batch_size, seq_len, num_heads, head_dim = shape_list(value)
+
+ tf.debugging.assert_equal(
+ seq_len % (window_overlap * 2), 0, message="Seq_len has to be multiple of 2 * window_overlap"
+ )
+ tf.debugging.assert_equal(
+ shape_list(attn_probs)[:3],
+ shape_list(value)[:3],
+ message="value and attn_probs must have same dims (except head_dim)",
+ )
+ tf.debugging.assert_equal(
+ shape_list(attn_probs)[3],
+ 2 * window_overlap + 1,
+ message="attn_probs last dim has to be 2 * window_overlap + 1",
+ )
+
+ chunks_count = seq_len // window_overlap - 1
+
+ # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size 2 window overlap
+ chunked_attn_probs = tf.reshape(
+ tf.transpose(attn_probs, (0, 2, 1, 3)),
+ (
+ batch_size * num_heads,
+ seq_len // window_overlap,
+ window_overlap,
+ 2 * window_overlap + 1,
+ ),
+ )
+
+ # group batch_size and num_heads dimensions into one
+ value = tf.reshape(
+ tf.transpose(value, (0, 2, 1, 3)),
+ (batch_size * num_heads, seq_len, head_dim),
+ )
+
+ # pad seq_len with w at the beginning of the sequence and another window overlap at the end
+ paddings = tf.convert_to_tensor([[0, 0], [window_overlap, window_overlap], [0, 0]])
+ padded_value = tf.pad(value, paddings, constant_values=-1)
+
+ # chunk padded_value into chunks of size 3 window overlap and an overlap of size window overlap
+ frame_size = 3 * window_overlap * head_dim
+ frame_hop_size = (shape_list(padded_value)[1] * head_dim - frame_size) // chunks_count
+ chunked_value = tf.signal.frame(
+ tf.reshape(padded_value, (batch_size * num_heads, -1)),
+ frame_size,
+ frame_hop_size,
+ )
+ chunked_value = tf.reshape(
+ chunked_value,
+ (batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim),
+ )
+
+ tf.debugging.assert_equal(
+ shape_list(chunked_value),
+ [batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim],
+ message="Chunked value has the wrong shape",
+ )
+
+ chunked_attn_probs = self._pad_and_diagonalize(chunked_attn_probs)
+ context = tf.einsum("bcwd,bcdh->bcwh", chunked_attn_probs, chunked_value)
+ context = tf.transpose(
+ tf.reshape(context, (batch_size, num_heads, seq_len, head_dim)),
+ (0, 2, 1, 3),
+ )
+
+ return context
+
+ @staticmethod
+ def _pad_and_transpose_last_two_dims(hidden_states_padded, paddings):
+ """pads rows and then flips rows and columns"""
+ hidden_states_padded = tf.pad(
+ hidden_states_padded, paddings
+ ) # padding value is not important because it will be overwritten
+ batch_size, chunk_size, seq_length, hidden_dim = shape_list(hidden_states_padded)
+ hidden_states_padded = tf.reshape(hidden_states_padded, (batch_size, chunk_size, hidden_dim, seq_length))
+
+ return hidden_states_padded
+
+ @staticmethod
+ def _pad_and_diagonalize(chunked_hidden_states):
+ """
+ shift every row 1 step right, converting columns into diagonals.
+
+ Example:
+
+ ```python
+ chunked_hidden_states: [
+ 0.4983,
+ 2.6918,
+ -0.0071,
+ 1.0492,
+ -1.8348,
+ 0.7672,
+ 0.2986,
+ 0.0285,
+ -0.7584,
+ 0.4206,
+ -0.0405,
+ 0.1599,
+ 2.0514,
+ -1.1600,
+ 0.5372,
+ 0.2629,
+ ]
+ window_overlap = num_rows = 4
+ ```
+
+ (pad & diagonalize) => [ 0.4983, 2.6918, -0.0071, 1.0492, 0.0000, 0.0000, 0.0000
+ 0.0000, -1.8348, 0.7672, 0.2986, 0.0285, 0.0000, 0.0000 0.0000, 0.0000, -0.7584, 0.4206,
+ -0.0405, 0.1599, 0.0000 0.0000, 0.0000, 0.0000, 2.0514, -1.1600, 0.5372, 0.2629 ]
+ """
+ total_num_heads, num_chunks, window_overlap, hidden_dim = shape_list(chunked_hidden_states)
+ paddings = tf.convert_to_tensor([[0, 0], [0, 0], [0, 0], [0, window_overlap + 1]])
+ chunked_hidden_states = tf.pad(
+ chunked_hidden_states, paddings
+ ) # total_num_heads x num_chunks x window_overlap x (hidden_dim+window_overlap+1). Padding value is not important because it'll be overwritten
+ chunked_hidden_states = tf.reshape(
+ chunked_hidden_states, (total_num_heads, num_chunks, -1)
+ ) # total_num_heads x num_chunks x window_overlapL+window_overlapwindow_overlap+window_overlap
+ chunked_hidden_states = chunked_hidden_states[
+ :, :, :-window_overlap
+ ] # total_num_heads x num_chunks x window_overlapL+window_overlapwindow_overlap
+ chunked_hidden_states = tf.reshape(
+ chunked_hidden_states,
+ (total_num_heads, num_chunks, window_overlap, window_overlap + hidden_dim),
+ ) # total_num_heads x num_chunks, window_overlap x hidden_dim+window_overlap
+ chunked_hidden_states = chunked_hidden_states[:, :, :, :-1]
+
+ return chunked_hidden_states
+
+ @staticmethod
+ def _chunk(hidden_states, window_overlap):
+ """convert into overlapping chunks. Chunk size = 2w, overlap size = w"""
+ batch_size, seq_length, hidden_dim = shape_list(hidden_states)
+ num_output_chunks = 2 * (seq_length // (2 * window_overlap)) - 1
+
+ # define frame size and frame stride (similar to convolution)
+ frame_hop_size = window_overlap * hidden_dim
+ frame_size = 2 * frame_hop_size
+ hidden_states = tf.reshape(hidden_states, (batch_size, seq_length * hidden_dim))
+
+ # chunk with overlap
+ chunked_hidden_states = tf.signal.frame(hidden_states, frame_size, frame_hop_size)
+
+ tf.debugging.assert_equal(
+ shape_list(chunked_hidden_states),
+ [batch_size, num_output_chunks, frame_size],
+ message=(
+ "Make sure chunking is correctly applied. `Chunked hidden states should have output dimension"
+ f" {[batch_size, frame_size, num_output_chunks]}, but got {shape_list(chunked_hidden_states)}."
+ ),
+ )
+
+ chunked_hidden_states = tf.reshape(
+ chunked_hidden_states,
+ (batch_size, num_output_chunks, 2 * window_overlap, hidden_dim),
+ )
+
+ return chunked_hidden_states
+
+ @staticmethod
+ def _get_global_attn_indices(is_index_global_attn):
+ """compute global attn indices required throughout forward pass"""
+ # helper variable
+ num_global_attn_indices = tf.math.count_nonzero(is_index_global_attn, axis=1)
+ num_global_attn_indices = tf.cast(num_global_attn_indices, dtype=tf.constant(1).dtype)
+
+ # max number of global attn indices in batch
+ max_num_global_attn_indices = tf.reduce_max(num_global_attn_indices)
+
+ # indices of global attn
+ is_index_global_attn_nonzero = tf.where(is_index_global_attn)
+
+ # helper variable
+ is_local_index_global_attn = tf.range(max_num_global_attn_indices) < tf.expand_dims(
+ num_global_attn_indices, axis=-1
+ )
+
+ # location of the non-padding values within global attention indices
+ is_local_index_global_attn_nonzero = tf.where(is_local_index_global_attn)
+
+ # location of the padding values within global attention indices
+ is_local_index_no_global_attn_nonzero = tf.where(tf.math.logical_not(is_local_index_global_attn))
+
+ return (
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ )
+
+ def _concat_with_global_key_attn_probs(
+ self,
+ attn_scores,
+ key_vectors,
+ query_vectors,
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ ):
+ batch_size = shape_list(key_vectors)[0]
+
+ # select global key vectors
+ global_key_vectors = tf.gather_nd(key_vectors, is_index_global_attn_nonzero)
+
+ # create only global key vectors
+ key_vectors_only_global = tf.scatter_nd(
+ is_local_index_global_attn_nonzero,
+ global_key_vectors,
+ shape=(
+ batch_size,
+ max_num_global_attn_indices,
+ self.num_heads,
+ self.head_dim,
+ ),
+ )
+
+ # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
+ attn_probs_from_global_key = tf.einsum("blhd,bshd->blhs", query_vectors, key_vectors_only_global)
+
+ # (batch_size, max_num_global_attn_indices, seq_len, num_heads)
+ attn_probs_from_global_key_trans = tf.transpose(attn_probs_from_global_key, (0, 3, 1, 2))
+ mask_shape = (shape_list(is_local_index_no_global_attn_nonzero)[0],) + tuple(
+ shape_list(attn_probs_from_global_key_trans)[-2:]
+ )
+ mask = tf.ones(mask_shape) * -10000.0
+ mask = tf.cast(mask, dtype=attn_probs_from_global_key_trans.dtype)
+
+ # scatter mask
+ attn_probs_from_global_key_trans = tf.tensor_scatter_nd_update(
+ attn_probs_from_global_key_trans,
+ is_local_index_no_global_attn_nonzero,
+ mask,
+ )
+
+ # (batch_size, seq_len, num_heads, max_num_global_attn_indices)
+ attn_probs_from_global_key = tf.transpose(attn_probs_from_global_key_trans, (0, 2, 3, 1))
+
+ # concat to attn_probs
+ # (batch_size, seq_len, num_heads, extra attention count + 2*window+1)
+ attn_scores = tf.concat((attn_probs_from_global_key, attn_scores), axis=-1)
+
+ return attn_scores
+
+ def _compute_attn_output_with_global_indices(
+ self,
+ value_vectors,
+ attn_probs,
+ max_num_global_attn_indices,
+ is_index_global_attn_nonzero,
+ is_local_index_global_attn_nonzero,
+ ):
+ batch_size = shape_list(attn_probs)[0]
+
+ # cut local attn probs to global only
+ attn_probs_only_global = attn_probs[:, :, :, :max_num_global_attn_indices]
+
+ # select global value vectors
+ global_value_vectors = tf.gather_nd(value_vectors, is_index_global_attn_nonzero)
+
+ # create only global value vectors
+ value_vectors_only_global = tf.scatter_nd(
+ is_local_index_global_attn_nonzero,
+ global_value_vectors,
+ shape=(
+ batch_size,
+ max_num_global_attn_indices,
+ self.num_heads,
+ self.head_dim,
+ ),
+ )
+
+ # compute attn output only global
+ attn_output_only_global = tf.einsum("blhs,bshd->blhd", attn_probs_only_global, value_vectors_only_global)
+
+ # reshape attn probs
+ attn_probs_without_global = attn_probs[:, :, :, max_num_global_attn_indices:]
+
+ # compute attn output with global
+ attn_output_without_global = self._sliding_chunks_matmul_attn_probs_value(
+ attn_probs_without_global, value_vectors, self.one_sided_attn_window_size
+ )
+
+ return attn_output_only_global + attn_output_without_global
+
+ def _compute_global_attn_output_from_hidden(
+ self,
+ attn_output,
+ hidden_states,
+ max_num_global_attn_indices,
+ layer_head_mask,
+ is_local_index_global_attn_nonzero,
+ is_index_global_attn_nonzero,
+ is_local_index_no_global_attn_nonzero,
+ is_index_masked,
+ training,
+ ):
+ batch_size, seq_len = shape_list(hidden_states)[:2]
+
+ # prepare global hidden states
+ global_attn_hidden_states = tf.gather_nd(hidden_states, is_index_global_attn_nonzero)
+ global_attn_hidden_states = tf.scatter_nd(
+ is_local_index_global_attn_nonzero,
+ global_attn_hidden_states,
+ shape=(batch_size, max_num_global_attn_indices, self.embed_dim),
+ )
+
+ # global key, query, value
+ global_query_vectors_only_global = self.query_global(global_attn_hidden_states)
+ global_key_vectors = self.key_global(hidden_states)
+ global_value_vectors = self.value_global(hidden_states)
+
+ # normalize
+ global_query_vectors_only_global /= tf.math.sqrt(
+ tf.cast(self.head_dim, dtype=global_query_vectors_only_global.dtype)
+ )
+ global_query_vectors_only_global = self.reshape_and_transpose(global_query_vectors_only_global, batch_size)
+ global_key_vectors = self.reshape_and_transpose(global_key_vectors, batch_size)
+ global_value_vectors = self.reshape_and_transpose(global_value_vectors, batch_size)
+
+ # compute attn scores
+ global_attn_scores = tf.matmul(global_query_vectors_only_global, global_key_vectors, transpose_b=True)
+
+ tf.debugging.assert_equal(
+ shape_list(global_attn_scores),
+ [batch_size * self.num_heads, max_num_global_attn_indices, seq_len],
+ message=(
+ "global_attn_scores have the wrong size. Size should be"
+ f" {(batch_size * self.num_heads, max_num_global_attn_indices, seq_len)}, but is"
+ f" {shape_list(global_attn_scores)}."
+ ),
+ )
+
+ global_attn_scores = tf.reshape(
+ global_attn_scores,
+ (batch_size, self.num_heads, max_num_global_attn_indices, seq_len),
+ )
+ global_attn_scores_trans = tf.transpose(global_attn_scores, (0, 2, 1, 3))
+ mask_shape = (shape_list(is_local_index_no_global_attn_nonzero)[0],) + tuple(
+ shape_list(global_attn_scores_trans)[-2:]
+ )
+ global_attn_mask = tf.ones(mask_shape) * -10000.0
+ global_attn_mask = tf.cast(global_attn_mask, dtype=global_attn_scores_trans.dtype)
+
+ # scatter mask
+ global_attn_scores_trans = tf.tensor_scatter_nd_update(
+ global_attn_scores_trans,
+ is_local_index_no_global_attn_nonzero,
+ global_attn_mask,
+ )
+ global_attn_scores = tf.transpose(global_attn_scores_trans, (0, 2, 1, 3))
+
+ # mask global attn scores
+ attn_mask = tf.tile(is_index_masked[:, None, None, :], (1, shape_list(global_attn_scores)[1], 1, 1))
+ global_attn_scores = tf.where(attn_mask, -10000.0, global_attn_scores)
+ global_attn_scores = tf.reshape(
+ global_attn_scores,
+ (batch_size * self.num_heads, max_num_global_attn_indices, seq_len),
+ )
+
+ # compute global attn probs
+ global_attn_probs_float = stable_softmax(global_attn_scores, axis=-1)
+
+ # apply layer head masking
+ if layer_head_mask is not None:
+ tf.debugging.assert_equal(
+ shape_list(layer_head_mask),
+ [self.num_heads],
+ message=(
+ f"Head mask for a single layer should be of size {(self.num_heads)}, but is"
+ f" {shape_list(layer_head_mask)}"
+ ),
+ )
+ global_attn_probs_float = tf.reshape(layer_head_mask, (1, -1, 1, 1)) * tf.reshape(
+ global_attn_probs_float, (batch_size, self.num_heads, max_num_global_attn_indices, seq_len)
+ )
+ global_attn_probs_float = tf.reshape(
+ global_attn_probs_float, (batch_size * self.num_heads, max_num_global_attn_indices, seq_len)
+ )
+
+ # dropout
+ global_attn_probs = self.global_dropout(global_attn_probs_float, training=training)
+
+ # global attn output
+ global_attn_output = tf.matmul(global_attn_probs, global_value_vectors)
+
+ tf.debugging.assert_equal(
+ shape_list(global_attn_output),
+ [batch_size * self.num_heads, max_num_global_attn_indices, self.head_dim],
+ message=(
+ "global_attn_output tensor has the wrong size. Size should be"
+ f" {(batch_size * self.num_heads, max_num_global_attn_indices, self.head_dim)}, but is"
+ f" {shape_list(global_attn_output)}."
+ ),
+ )
+
+ global_attn_output = tf.reshape(
+ global_attn_output,
+ (batch_size, self.num_heads, max_num_global_attn_indices, self.head_dim),
+ )
+
+ # get only non zero global attn output
+ nonzero_global_attn_output = tf.gather_nd(
+ tf.transpose(global_attn_output, (0, 2, 1, 3)),
+ is_local_index_global_attn_nonzero,
+ )
+ nonzero_global_attn_output = tf.reshape(
+ nonzero_global_attn_output,
+ (shape_list(is_local_index_global_attn_nonzero)[0], -1),
+ )
+
+ # overwrite values with global attention
+ attn_output = tf.tensor_scatter_nd_update(
+ attn_output, is_index_global_attn_nonzero, nonzero_global_attn_output
+ )
+
+ global_attn_probs = tf.reshape(
+ global_attn_probs, (batch_size, self.num_heads, max_num_global_attn_indices, seq_len)
+ )
+
+ return attn_output, global_attn_probs
+
+ def reshape_and_transpose(self, vector, batch_size):
+ return tf.reshape(
+ tf.transpose(
+ tf.reshape(vector, (batch_size, -1, self.num_heads, self.head_dim)),
+ (0, 2, 1, 3),
+ ),
+ (batch_size * self.num_heads, -1, self.head_dim),
+ )
+
+
+class TFLongformerAttention(keras.layers.Layer):
+ def __init__(self, config, layer_id=0, **kwargs):
+ super().__init__(**kwargs)
+
+ self.self_attention = TFLongformerSelfAttention(config, layer_id, name="self")
+ self.dense_output = TFLongformerSelfOutput(config, name="output")
+
+ def prune_heads(self, heads):
+ raise NotImplementedError
+
+ def call(self, inputs, training=False):
+ (
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ is_index_masked,
+ is_index_global_attn,
+ is_global_attn,
+ ) = inputs
+
+ self_outputs = self.self_attention(
+ [hidden_states, attention_mask, layer_head_mask, is_index_masked, is_index_global_attn, is_global_attn],
+ training=training,
+ )
+ attention_output = self.dense_output(self_outputs[0], hidden_states, training=training)
+ 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)
+
+
+class TFLongformerLayer(keras.layers.Layer):
+ def __init__(self, config, layer_id=0, **kwargs):
+ super().__init__(**kwargs)
+
+ self.attention = TFLongformerAttention(config, layer_id, name="attention")
+ self.intermediate = TFLongformerIntermediate(config, name="intermediate")
+ self.longformer_output = TFLongformerOutput(config, name="output")
+
+ def call(self, inputs, training=False):
+ (
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ is_index_masked,
+ is_index_global_attn,
+ is_global_attn,
+ ) = inputs
+
+ attention_outputs = self.attention(
+ [hidden_states, attention_mask, layer_head_mask, is_index_masked, is_index_global_attn, is_global_attn],
+ training=training,
+ )
+ attention_output = attention_outputs[0]
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.longformer_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, "longformer_output", None) is not None:
+ with tf.name_scope(self.longformer_output.name):
+ self.longformer_output.build(None)
+
+
+class TFLongformerEncoder(keras.layers.Layer):
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+
+ self.output_hidden_states = config.output_hidden_states
+ self.output_attentions = config.output_attentions
+ self.layer = [TFLongformerLayer(config, i, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
+
+ def call(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ padding_len=0,
+ is_index_masked=None,
+ is_index_global_attn=None,
+ is_global_attn=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ training=False,
+ ):
+ all_hidden_states = () if output_hidden_states else None
+ all_attentions = all_global_attentions = () if output_attentions else None
+
+ for idx, layer_module in enumerate(self.layer):
+ if output_hidden_states:
+ hidden_states_to_add = hidden_states[:, :-padding_len] if padding_len > 0 else hidden_states
+ all_hidden_states = all_hidden_states + (hidden_states_to_add,)
+
+ layer_outputs = layer_module(
+ [
+ hidden_states,
+ attention_mask,
+ head_mask[idx] if head_mask is not None else None,
+ is_index_masked,
+ is_index_global_attn,
+ is_global_attn,
+ ],
+ training=training,
+ )
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ # bzs x seq_len x num_attn_heads x (num_global_attn + attention_window_len + 1) => bzs x num_attn_heads x seq_len x (num_global_attn + attention_window_len + 1)
+ all_attentions = all_attentions + (tf.transpose(layer_outputs[1], (0, 2, 1, 3)),)
+
+ # bzs x num_attn_heads x num_global_attn x seq_len => bzs x num_attn_heads x seq_len x num_global_attn
+ all_global_attentions = all_global_attentions + (tf.transpose(layer_outputs[2], (0, 1, 3, 2)),)
+
+ # Add last layer
+ if output_hidden_states:
+ hidden_states_to_add = hidden_states[:, :-padding_len] if padding_len > 0 else hidden_states
+ all_hidden_states = all_hidden_states + (hidden_states_to_add,)
+
+ # undo padding
+ # unpad `hidden_states` because the calling function is expecting a length == input_ids.size(1)
+ hidden_states = hidden_states[:, :-padding_len] if padding_len > 0 else hidden_states
+ if output_attentions:
+ all_attentions = (
+ tuple([state[:, :, :-padding_len, :] for state in all_attentions])
+ if padding_len > 0
+ else all_attentions
+ )
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, all_hidden_states, all_attentions, all_global_attentions] if v is not None
+ )
+
+ return TFLongformerBaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_attentions,
+ global_attentions=all_global_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 TFLongformerMainLayer(keras.layers.Layer):
+ config_class = LongformerConfig
+
+ def __init__(self, config, add_pooling_layer=True, **kwargs):
+ super().__init__(**kwargs)
+
+ if isinstance(config.attention_window, int):
+ assert config.attention_window % 2 == 0, "`config.attention_window` has to be an even value"
+ assert config.attention_window > 0, "`config.attention_window` has to be positive"
+ config.attention_window = [config.attention_window] * config.num_hidden_layers # one value per layer
+ else:
+ assert len(config.attention_window) == config.num_hidden_layers, (
+ "`len(config.attention_window)` should equal `config.num_hidden_layers`. "
+ f"Expected {config.num_hidden_layers}, given {len(config.attention_window)}"
+ )
+
+ self.config = config
+ self.num_hidden_layers = config.num_hidden_layers
+ self.initializer_range = config.initializer_range
+ self.output_attentions = config.output_attentions
+ self.output_hidden_states = config.output_hidden_states
+ self.return_dict = config.use_return_dict
+ self.pad_token_id = config.pad_token_id
+ self.attention_window = config.attention_window
+ self.embeddings = TFLongformerEmbeddings(config, name="embeddings")
+ self.encoder = TFLongformerEncoder(config, name="encoder")
+ self.pooler = TFLongformerPooler(config, name="pooler") if add_pooling_layer else None
+
+ 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
+
+ @unpack_inputs
+ def call(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ head_mask=None,
+ global_attention_mask=None,
+ token_type_ids=None,
+ position_ids=None,
+ inputs_embeds=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ training=False,
+ ):
+ if input_ids is not None and not isinstance(input_ids, tf.Tensor):
+ input_ids = tf.convert_to_tensor(input_ids, dtype=tf.int64)
+ elif input_ids is not None:
+ input_ids = tf.cast(input_ids, tf.int64)
+
+ if attention_mask is not None and not isinstance(attention_mask, tf.Tensor):
+ attention_mask = tf.convert_to_tensor(attention_mask, dtype=tf.int64)
+ elif attention_mask is not None:
+ attention_mask = tf.cast(attention_mask, tf.int64)
+
+ if global_attention_mask is not None and not isinstance(global_attention_mask, tf.Tensor):
+ global_attention_mask = tf.convert_to_tensor(global_attention_mask, dtype=tf.int64)
+ elif global_attention_mask is not None:
+ global_attention_mask = tf.cast(global_attention_mask, tf.int64)
+
+ 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.cast(tf.fill(input_shape, 1), tf.int64)
+
+ if token_type_ids is None:
+ token_type_ids = tf.cast(tf.fill(input_shape, 0), tf.int64)
+
+ # merge `global_attention_mask` and `attention_mask`
+ if global_attention_mask is not None:
+ attention_mask = self._merge_to_attention_mask(attention_mask, global_attention_mask)
+
+ (
+ padding_len,
+ input_ids,
+ attention_mask,
+ token_type_ids,
+ position_ids,
+ inputs_embeds,
+ ) = self._pad_to_window_size(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ pad_token_id=self.pad_token_id,
+ )
+
+ # is index masked or global attention
+ is_index_masked = tf.math.less(attention_mask, 1)
+ is_index_global_attn = tf.math.greater(attention_mask, 1)
+ is_global_attn = tf.math.reduce_any(is_index_global_attn)
+
+ # We create a 3D attention mask from a 2D tensor mask.
+ # Sizes are [batch_size, to_seq_length, 1, 1]
+ # 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)
+ extended_attention_mask = tf.reshape(attention_mask, (attention_mask_shape[0], attention_mask_shape[1], 1, 1))
+
+ # Since attention_mask is 1.0 for positions we want to attend locally and 0.0 for
+ # masked and global attn 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(tf.math.abs(1 - extended_attention_mask), tf.dtypes.float32) * -10000.0
+ embedding_output = self.embeddings(
+ input_ids,
+ position_ids,
+ token_type_ids,
+ inputs_embeds,
+ training=training,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ padding_len=padding_len,
+ is_index_masked=is_index_masked,
+ is_index_global_attn=is_index_global_attn,
+ is_global_attn=is_global_attn,
+ 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(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (
+ sequence_output,
+ pooled_output,
+ ) + encoder_outputs[1:]
+
+ return TFLongformerBaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ global_attentions=encoder_outputs.global_attentions,
+ )
+
+ def _pad_to_window_size(
+ self,
+ input_ids,
+ attention_mask,
+ token_type_ids,
+ position_ids,
+ inputs_embeds,
+ pad_token_id,
+ ):
+ """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
+ # padding
+ attention_window = (
+ self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
+ )
+
+ assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"
+
+ input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
+ batch_size, seq_len = input_shape[:2]
+ padding_len = (attention_window - seq_len % attention_window) % attention_window
+
+ paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])
+
+ if input_ids is not None:
+ input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)
+
+ if position_ids is not None:
+ # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
+ position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)
+
+ if inputs_embeds is not None:
+ if padding_len > 0:
+ input_ids_padding = tf.cast(tf.fill((batch_size, padding_len), self.pad_token_id), tf.int64)
+ inputs_embeds_padding = self.embeddings(input_ids_padding)
+ inputs_embeds = tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)
+
+ attention_mask = tf.pad(attention_mask, paddings, constant_values=False) # no attention on the padding tokens
+ token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0) # pad with token_type_id = 0
+
+ return (
+ padding_len,
+ input_ids,
+ attention_mask,
+ token_type_ids,
+ position_ids,
+ inputs_embeds,
+ )
+
+ @staticmethod
+ def _merge_to_attention_mask(attention_mask: tf.Tensor, global_attention_mask: tf.Tensor):
+ # longformer self attention expects attention mask to have 0 (no attn), 1 (local attn), 2 (global attn)
+ # (global_attention_mask + 1) => 1 for local attention, 2 for global attention
+ # => final attention_mask => 0 for no attention, 1 for local attention 2 for global attention
+ if attention_mask is not None:
+ attention_mask = attention_mask * (global_attention_mask + 1)
+ else:
+ # simply use `global_attention_mask` as `attention_mask`
+ # if no `attention_mask` is given
+ attention_mask = global_attention_mask + 1
+
+ return attention_mask
+
+ 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)
+
+
+class TFLongformerPreTrainedModel(TFPreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = LongformerConfig
+ base_model_prefix = "longformer"
+
+ @property
+ def input_signature(self):
+ sig = super().input_signature
+ sig["global_attention_mask"] = tf.TensorSpec((None, None), tf.int32, name="global_attention_mask")
+ return sig
+
+
+LONGFORMER_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 ([`LongformerConfig`]): 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.
+"""
+
+
+LONGFORMER_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`np.ndarray` 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 (`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)
+ head_mask (`np.ndarray` or `tf.Tensor` of shape `(encoder_layers, encoder_attention_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**.
+
+ global_attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
+ Mask to decide the attention given on each token, local attention or global attention. Tokens with global
+ attention attends to all other tokens, and all other tokens attend to them. This is important for
+ task-specific finetuning because it makes the model more flexible at representing the task. For example,
+ for classification, the token should be given global attention. For QA, all question tokens should also
+ have global attention. Please refer to the [Longformer paper](https://arxiv.org/abs/2004.05150) for more
+ details. Mask values selected in `[0, 1]`:
+
+ - 0 for local attention (a sliding window attention),
+ - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them).
+
+ 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. 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 Longformer Model outputting raw hidden-states without any specific head on top.",
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerModel(TFLongformerPreTrainedModel):
+ """
+
+ This class copies code from [`TFRobertaModel`] and overwrites standard self-attention with longformer
+ self-attention to provide the ability to process long sequences following the self-attention approach described in
+ [Longformer: the Long-Document Transformer](https://arxiv.org/abs/2004.05150) by Iz Beltagy, Matthew E. Peters, and
+ Arman Cohan. Longformer self-attention combines a local (sliding window) and global attention to extend to long
+ documents without the O(n^2) increase in memory and compute.
+
+ The self-attention module `TFLongformerSelfAttention` implemented here supports the combination of local and global
+ attention but it lacks support for autoregressive attention and dilated attention. Autoregressive and dilated
+ attention are more relevant for autoregressive language modeling than finetuning on downstream tasks. Future
+ release will add support for autoregressive attention, but the support for dilated attention requires a custom CUDA
+ kernel to be memory and compute efficient.
+
+ """
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.longformer = TFLongformerMainLayer(config, name="longformer")
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ global_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[TFLongformerBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
+ outputs = self.longformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=global_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, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.build(None)
+
+
+@add_start_docstrings(
+ """Longformer Model with a `language modeling` head on top.""",
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerForMaskedLM(TFLongformerPreTrainedModel, TFMaskedLanguageModelingLoss):
+ # 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"pooler"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.longformer = TFLongformerMainLayer(config, add_pooling_layer=False, name="longformer")
+ self.lm_head = TFLongformerLMHead(config, self.longformer.embeddings, name="lm_head")
+
+ def get_lm_head(self):
+ return self.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.lm_head.name
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint="allenai/longformer-base-4096",
+ output_type=TFLongformerMaskedLMOutput,
+ config_class=_CONFIG_FOR_DOC,
+ mask="",
+ expected_output="' Paris'",
+ expected_loss=0.44,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ global_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[TFLongformerMaskedLMOutput, 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]`
+ """
+
+ outputs = self.longformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=global_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.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[2:]
+
+ return ((loss,) + output) if loss is not None else output
+
+ return TFLongformerMaskedLMOutput(
+ loss=loss,
+ logits=prediction_scores,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.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(
+ """
+ Longformer Model with a span classification head on top for extractive question-answering tasks like SQuAD /
+ TriviaQA (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerForQuestionAnswering(TFLongformerPreTrainedModel, TFQuestionAnsweringLoss):
+ # 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"pooler"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+ self.longformer = TFLongformerMainLayer(config, add_pooling_layer=False, name="longformer")
+ 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(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint="allenai/longformer-large-4096-finetuned-triviaqa",
+ output_type=TFLongformerQuestionAnsweringModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output="' puppet'",
+ expected_loss=0.96,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ global_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[TFLongformerQuestionAnsweringModelOutput, 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.
+ """
+
+ if input_ids is not None and not isinstance(input_ids, tf.Tensor):
+ input_ids = tf.convert_to_tensor(input_ids, dtype=tf.int64)
+ elif input_ids is not None:
+ input_ids = tf.cast(input_ids, tf.int64)
+
+ if attention_mask is not None and not isinstance(attention_mask, tf.Tensor):
+ attention_mask = tf.convert_to_tensor(attention_mask, dtype=tf.int64)
+ elif attention_mask is not None:
+ attention_mask = tf.cast(attention_mask, tf.int64)
+
+ if global_attention_mask is not None and not isinstance(global_attention_mask, tf.Tensor):
+ global_attention_mask = tf.convert_to_tensor(global_attention_mask, dtype=tf.int64)
+ elif global_attention_mask is not None:
+ global_attention_mask = tf.cast(global_attention_mask, tf.int64)
+
+ # set global attention on question tokens
+ if global_attention_mask is None and input_ids is not None:
+ if shape_list(tf.where(input_ids == self.config.sep_token_id))[0] != 3 * shape_list(input_ids)[0]:
+ logger.warning(
+ f"There should be exactly three separator tokens: {self.config.sep_token_id} in every sample for"
+ " questions answering. You might also consider to set `global_attention_mask` manually in the"
+ " forward function to avoid this. This is most likely an error. The global attention is disabled"
+ " for this forward pass."
+ )
+ global_attention_mask = tf.cast(tf.fill(shape_list(input_ids), value=0), tf.int64)
+ else:
+ logger.warning_once("Initializing global attention on question tokens...")
+ # put global attention on all tokens until `config.sep_token_id` is reached
+ sep_token_indices = tf.where(input_ids == self.config.sep_token_id)
+ sep_token_indices = tf.cast(sep_token_indices, dtype=tf.int64)
+ global_attention_mask = _compute_global_attention_mask(shape_list(input_ids), sep_token_indices)
+
+ outputs = self.longformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=global_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(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[2:]
+
+ return ((loss,) + output) if loss is not None else output
+
+ return TFLongformerQuestionAnsweringModelOutput(
+ loss=loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.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])
+
+
+class TFLongformerClassificationHead(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),
+ 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), name="out_proj"
+ )
+ self.config = config
+
+ def call(self, hidden_states, training=False):
+ hidden_states = hidden_states[:, 0, :] # take token (equiv. to [CLS])
+ hidden_states = self.dropout(hidden_states, training=training)
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states, training=training)
+ output = self.out_proj(hidden_states)
+ return 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])
+ 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(
+ """
+ Longformer Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+ pooled output) e.g. for GLUE tasks.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerForSequenceClassification(TFLongformerPreTrainedModel, TFSequenceClassificationLoss):
+ # 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"pooler"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+
+ self.longformer = TFLongformerMainLayer(config, add_pooling_layer=False, name="longformer")
+ self.classifier = TFLongformerClassificationHead(config, name="classifier")
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFLongformerSequenceClassifierOutput,
+ 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,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ position_ids: np.ndarray | tf.Tensor | None = None,
+ global_attention_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[TFLongformerSequenceClassifierOutput, Tuple[tf.Tensor]]:
+ if input_ids is not None and not isinstance(input_ids, tf.Tensor):
+ input_ids = tf.convert_to_tensor(input_ids, dtype=tf.int64)
+ elif input_ids is not None:
+ input_ids = tf.cast(input_ids, tf.int64)
+
+ if attention_mask is not None and not isinstance(attention_mask, tf.Tensor):
+ attention_mask = tf.convert_to_tensor(attention_mask, dtype=tf.int64)
+ elif attention_mask is not None:
+ attention_mask = tf.cast(attention_mask, tf.int64)
+
+ if global_attention_mask is not None and not isinstance(global_attention_mask, tf.Tensor):
+ global_attention_mask = tf.convert_to_tensor(global_attention_mask, dtype=tf.int64)
+ elif global_attention_mask is not None:
+ global_attention_mask = tf.cast(global_attention_mask, tf.int64)
+
+ if global_attention_mask is None and input_ids is not None:
+ logger.warning_once("Initializing global attention on CLS token...")
+ # global attention on cls token
+ global_attention_mask = tf.zeros_like(input_ids)
+ updates = tf.ones(shape_list(input_ids)[0], dtype=tf.int64)
+ indices = tf.pad(
+ tensor=tf.expand_dims(tf.range(shape_list(input_ids)[0], dtype=tf.int64), axis=1),
+ paddings=[[0, 0], [0, 1]],
+ constant_values=0,
+ )
+ global_attention_mask = tf.tensor_scatter_nd_update(
+ global_attention_mask,
+ indices,
+ updates,
+ )
+
+ outputs = self.longformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=global_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.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 TFLongformerSequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.build(None)
+ if getattr(self, "classifier", None) is not None:
+ with tf.name_scope(self.classifier.name):
+ self.classifier.build(None)
+
+
+@add_start_docstrings(
+ """
+ Longformer 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.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerForMultipleChoice(TFLongformerPreTrainedModel, 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_missing = [r"dropout"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.longformer = TFLongformerMainLayer(config, name="longformer")
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
+ self.classifier = keras.layers.Dense(
+ 1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
+ )
+ self.config = config
+
+ @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"),
+ "global_attention_mask": tf.TensorSpec((None, None, None), tf.int32, name="global_attention_mask"),
+ }
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(
+ LONGFORMER_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
+ )
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFLongformerMultipleChoiceModelOutput,
+ 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,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ position_ids: np.ndarray | tf.Tensor | None = None,
+ global_attention_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[TFLongformerMultipleChoiceModelOutput, 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_global_attention_mask = (
+ tf.reshape(global_attention_mask, (-1, shape_list(global_attention_mask)[-1]))
+ if global_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
+ )
+
+ outputs = self.longformer(
+ flat_input_ids,
+ position_ids=flat_position_ids,
+ token_type_ids=flat_token_type_ids,
+ attention_mask=flat_attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=flat_global_attention_mask,
+ inputs_embeds=flat_inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ training=training,
+ )
+ pooled_output = outputs[1]
+
+ pooled_output = self.dropout(pooled_output)
+ 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,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TFLongformerMultipleChoiceModelOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.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(
+ """
+ Longformer 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.
+ """,
+ LONGFORMER_START_DOCSTRING,
+)
+class TFLongformerForTokenClassification(TFLongformerPreTrainedModel, TFTokenClassificationLoss):
+ # 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"pooler"]
+ _keys_to_ignore_on_load_missing = [r"dropout"]
+
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+ self.longformer = TFLongformerMainLayer(config=config, add_pooling_layer=False, name="longformer")
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
+ 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(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFLongformerTokenClassifierOutput,
+ 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,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ position_ids: np.ndarray | tf.Tensor | None = None,
+ global_attention_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: Optional[Union[np.array, tf.Tensor]] = None,
+ training: Optional[bool] = False,
+ ) -> Union[TFLongformerTokenClassifierOutput, 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.longformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ head_mask=head_mask,
+ global_attention_mask=global_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)
+ 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 TFLongformerTokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ global_attentions=outputs.global_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "longformer", None) is not None:
+ with tf.name_scope(self.longformer.name):
+ self.longformer.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])
+
+
+__all__ = [
+ "TFLongformerForMaskedLM",
+ "TFLongformerForMultipleChoice",
+ "TFLongformerForQuestionAnswering",
+ "TFLongformerForSequenceClassification",
+ "TFLongformerForTokenClassification",
+ "TFLongformerModel",
+ "TFLongformerPreTrainedModel",
+ "TFLongformerSelfAttention",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer.py b/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..afecf750135b0d4a45f79ea91cfb223beefa814c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer.py
@@ -0,0 +1,402 @@
+# coding=utf-8
+# Copyright 2020 The Allen Institute for AI team 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
+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()
+# Copied from transformers.models.roberta.tokenization_roberta.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.roberta.tokenization_roberta.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.roberta.tokenization_roberta.RobertaTokenizer with FacebookAI/roberta-base->allenai/longformer-base-4096, RoBERTa->Longformer all-casing, RobertaTokenizer->LongformerTokenizer
+class LongformerTokenizer(PreTrainedTokenizer):
+ """
+ Constructs a Longformer tokenizer, derived from the GPT-2 tokenizer, using 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 LongformerTokenizer
+
+ >>> tokenizer = LongformerTokenizer.from_pretrained("allenai/longformer-base-4096")
+ >>> tokenizer("Hello world")["input_ids"]
+ [0, 31414, 232, 2]
+
+ >>> tokenizer(" Hello world")["input_ids"]
+ [0, 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 `""`):
+ 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. (Longformer tokenizer detect beginning of words by the preceding space).
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+
+ def __init__(
+ self,
+ vocab_file,
+ merges_file,
+ errors="replace",
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ add_prefix_space=False,
+ **kwargs,
+ ):
+ bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
+ pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_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
+ 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
+
+ # Mask token behave like a normal word, i.e. include the space before it
+ mask_token = (
+ AddedToken(mask_token, lstrip=True, rstrip=False, normalized=False)
+ if isinstance(mask_token, str)
+ else mask_token
+ )
+
+ # these special tokens are not part of the vocab.json, let's add them in the correct order
+
+ 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,
+ **kwargs,
+ )
+
+ @property
+ def vocab_size(self):
+ return len(self.encoder)
+
+ def get_vocab(self):
+ vocab = dict(self.encoder).copy()
+ vocab.update(self.added_tokens_encoder)
+ return vocab
+
+ 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
+
+ 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 Longformer 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]:
+ """
+ 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 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]:
+ """
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. Longformer does not
+ make use of token type ids, therefore a list of zeros is returned.
+
+ 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 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)
+
+
+__all__ = ["LongformerTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer_fast.py b/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..b8111b3d8a25f35bb1fae0f1aa28e28b35fa3cfd
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/longformer/tokenization_longformer_fast.py
@@ -0,0 +1,265 @@
+# coding=utf-8
+# Copyright 2020 The Allen Institute for AI team 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 classes for Longformer."""
+
+import json
+from typing import List, Optional, Tuple
+
+from tokenizers import processors
+
+from ...tokenization_utils_base import AddedToken, BatchEncoding
+from ...tokenization_utils_fast import PreTrainedTokenizerFast
+from ...utils import logging
+from .tokenization_longformer import LongformerTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
+
+
+# Copied from transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast with FacebookAI/roberta-base->allenai/longformer-base-4096, RoBERTa->Longformer all-casing, Roberta->Longformer
+class LongformerTokenizerFast(PreTrainedTokenizerFast):
+ """
+ Construct a "fast" Longformer tokenizer (backed by HuggingFace's *tokenizers* library), derived from the GPT-2
+ tokenizer, using 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 LongformerTokenizerFast
+
+ >>> tokenizer = LongformerTokenizerFast.from_pretrained("allenai/longformer-base-4096")
+ >>> tokenizer("Hello world")["input_ids"]
+ [0, 31414, 232, 2]
+
+ >>> tokenizer(" Hello world")["input_ids"]
+ [0, 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 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.
+ 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. (Longformer tokenizer detect beginning of words by the preceding space).
+ trim_offsets (`bool`, *optional*, defaults to `True`):
+ Whether the post processing step should trim offsets to avoid including whitespaces.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ slow_tokenizer_class = LongformerTokenizer
+
+ def __init__(
+ self,
+ vocab_file=None,
+ merges_file=None,
+ tokenizer_file=None,
+ errors="replace",
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ add_prefix_space=False,
+ trim_offsets=True,
+ **kwargs,
+ ):
+ mask_token = (
+ AddedToken(mask_token, lstrip=True, rstrip=False, normalized=False)
+ if isinstance(mask_token, str)
+ else mask_token
+ )
+ super().__init__(
+ vocab_file,
+ merges_file,
+ tokenizer_file=tokenizer_file,
+ errors=errors,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ cls_token=cls_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ mask_token=mask_token,
+ add_prefix_space=add_prefix_space,
+ trim_offsets=trim_offsets,
+ **kwargs,
+ )
+
+ tokenizer_component = "post_processor"
+ tokenizer_component_instance = getattr(self.backend_tokenizer, tokenizer_component, None)
+ if tokenizer_component_instance:
+ state = json.loads(tokenizer_component_instance.__getstate__())
+
+ # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
+ if "sep" in state:
+ state["sep"] = tuple(state["sep"])
+ if "cls" in state:
+ state["cls"] = tuple(state["cls"])
+
+ changes_to_apply = False
+
+ if state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
+ state["add_prefix_space"] = add_prefix_space
+ changes_to_apply = True
+
+ if state.get("trim_offsets", trim_offsets) != trim_offsets:
+ state["trim_offsets"] = trim_offsets
+ changes_to_apply = True
+
+ if changes_to_apply:
+ component_class = getattr(processors, state.pop("type"))
+ new_value = component_class(**state)
+ setattr(self.backend_tokenizer, tokenizer_component, new_value)
+
+ @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.
+
+ Longformer tokenizer has a special mask token to be usable in the fill-mask pipeline. The mask token will greedily
+ comprise the space before the **.
+ """
+ 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.
+
+ This is needed to preserve backward compatibility with all the previously used models based on Longformer.
+ """
+ # 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 _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)
+
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+ output = [self.bos_token_id] + token_ids_0 + [self.eos_token_id]
+ if token_ids_1 is None:
+ return output
+
+ return output + [self.eos_token_id] + token_ids_1 + [self.eos_token_id]
+
+ 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. Longformer does not
+ make use of token type ids, therefore a list of zeros is returned.
+
+ 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 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]
+
+
+__all__ = ["LongformerTokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__init__.py b/lib/python3.12/site-packages/transformers/models/maskformer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..144ae7d0a7bab663d7903bb2dd4ccddf8cd63e35
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/__init__.py
@@ -0,0 +1,31 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_maskformer import *
+ from .configuration_maskformer_swin import *
+ from .feature_extraction_maskformer import *
+ from .image_processing_maskformer import *
+ from .modeling_maskformer import *
+ from .modeling_maskformer_swin import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e4475f6751e717355f26eba9f7a638314540b396
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e3b41515b26cf19eb750efb10fb47c942c529818
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer_swin.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer_swin.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e6234d5dbec683ee9a56784cbb99d75faa6ce535
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/configuration_maskformer_swin.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/feature_extraction_maskformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/feature_extraction_maskformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f9a6d25e61d3be44bfe9fd5252a7d17e4adb0ac6
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/feature_extraction_maskformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/image_processing_maskformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/image_processing_maskformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..171d8aee3e1ce80a29dd56521a257d20fa6349bf
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/image_processing_maskformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8b40d3092d1ce2bfd1bf47ffb9272509502e4e0d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer_swin.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer_swin.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..95200062d9d2249fea9fefe1d202c2fb65ca24aa
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/maskformer/__pycache__/modeling_maskformer_swin.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer.py b/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..8149e028b91c325d0342a7664b72adbf592ec836
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer.py
@@ -0,0 +1,226 @@
+# 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.
+"""MaskFormer model configuration"""
+
+from typing import Dict, Optional
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ...utils.backbone_utils import verify_backbone_config_arguments
+from ..auto import CONFIG_MAPPING
+from ..detr import DetrConfig
+from ..swin import SwinConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class MaskFormerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MaskFormerModel`]. It is used to instantiate a
+ MaskFormer 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 MaskFormer
+ [facebook/maskformer-swin-base-ade](https://huggingface.co/facebook/maskformer-swin-base-ade) architecture trained
+ on [ADE20k-150](https://huggingface.co/datasets/scene_parse_150).
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Currently, MaskFormer only supports the [Swin Transformer](swin) as backbone.
+
+ Args:
+ mask_feature_size (`int`, *optional*, defaults to 256):
+ The masks' features size, this value will also be used to specify the Feature Pyramid Network features'
+ size.
+ no_object_weight (`float`, *optional*, defaults to 0.1):
+ Weight to apply to the null (no object) class.
+ use_auxiliary_loss(`bool`, *optional*, defaults to `False`):
+ If `True` [`MaskFormerForInstanceSegmentationOutput`] will contain the auxiliary losses computed using the
+ logits from each decoder's stage.
+ backbone_config (`Dict`, *optional*):
+ The configuration passed to the backbone, if unset, the configuration corresponding to
+ `swin-base-patch4-window12-384` will be used.
+ 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.
+ decoder_config (`Dict`, *optional*):
+ The configuration passed to the transformer decoder model, if unset the base config for `detr-resnet-50`
+ will be used.
+ 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.
+ dice_weight (`float`, *optional*, defaults to 1.0):
+ The weight for the dice loss.
+ cross_entropy_weight (`float`, *optional*, defaults to 1.0):
+ The weight for the cross entropy loss.
+ mask_weight (`float`, *optional*, defaults to 20.0):
+ The weight for the mask loss.
+ output_auxiliary_logits (`bool`, *optional*):
+ Should the model output its `auxiliary_logits` or not.
+
+ Raises:
+ `ValueError`:
+ Raised if the backbone model type selected is not in `["swin"]` or the decoder model type selected is not
+ in `["detr"]`
+
+ Examples:
+
+ ```python
+ >>> from transformers import MaskFormerConfig, MaskFormerModel
+
+ >>> # Initializing a MaskFormer facebook/maskformer-swin-base-ade configuration
+ >>> configuration = MaskFormerConfig()
+
+ >>> # Initializing a model (with random weights) from the facebook/maskformer-swin-base-ade style configuration
+ >>> model = MaskFormerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+
+ """
+
+ model_type = "maskformer"
+ attribute_map = {"hidden_size": "mask_feature_size"}
+ backbones_supported = ["resnet", "swin"]
+ decoders_supported = ["detr"]
+
+ def __init__(
+ self,
+ fpn_feature_size: int = 256,
+ mask_feature_size: int = 256,
+ no_object_weight: float = 0.1,
+ use_auxiliary_loss: bool = False,
+ backbone_config: Optional[Dict] = None,
+ decoder_config: Optional[Dict] = None,
+ init_std: float = 0.02,
+ init_xavier_std: float = 1.0,
+ dice_weight: float = 1.0,
+ cross_entropy_weight: float = 1.0,
+ mask_weight: float = 20.0,
+ output_auxiliary_logits: Optional[bool] = None,
+ backbone: Optional[str] = None,
+ use_pretrained_backbone: bool = False,
+ use_timm_backbone: bool = False,
+ backbone_kwargs: Optional[Dict] = None,
+ **kwargs,
+ ):
+ if backbone_config is None and backbone is None:
+ # fall back to https://huggingface.co/microsoft/swin-base-patch4-window12-384-in22k
+ backbone_config = SwinConfig(
+ image_size=384,
+ num_channels=3,
+ patch_size=4,
+ embed_dim=128,
+ depths=[2, 2, 18, 2],
+ num_heads=[4, 8, 16, 32],
+ window_size=12,
+ drop_path_rate=0.3,
+ out_features=["stage1", "stage2", "stage3", "stage4"],
+ )
+ elif 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)
+
+ verify_backbone_config_arguments(
+ use_timm_backbone=use_timm_backbone,
+ use_pretrained_backbone=use_pretrained_backbone,
+ backbone=backbone,
+ backbone_config=backbone_config,
+ backbone_kwargs=backbone_kwargs,
+ )
+ # verify that the backbone is supported
+ if backbone_config is not None and backbone_config.model_type not in self.backbones_supported:
+ logger.warning_once(
+ f"Backbone {backbone_config.model_type} is not a supported model and may not be compatible with MaskFormer. "
+ f"Supported model types: {','.join(self.backbones_supported)}"
+ )
+
+ if decoder_config is None:
+ # fall back to https://huggingface.co/facebook/detr-resnet-50
+ decoder_config = DetrConfig()
+ else:
+ # verify that the decoder is supported
+ decoder_type = (
+ decoder_config.pop("model_type") if isinstance(decoder_config, dict) else decoder_config.model_type
+ )
+ if decoder_type not in self.decoders_supported:
+ raise ValueError(
+ f"Transformer Decoder {decoder_type} not supported, please use one of"
+ f" {','.join(self.decoders_supported)}"
+ )
+ if isinstance(decoder_config, dict):
+ config_class = CONFIG_MAPPING[decoder_type]
+ decoder_config = config_class.from_dict(decoder_config)
+
+ self.backbone_config = backbone_config
+ self.decoder_config = decoder_config
+ # main feature dimension for the model
+ self.fpn_feature_size = fpn_feature_size
+ self.mask_feature_size = mask_feature_size
+ # initializer
+ self.init_std = init_std
+ self.init_xavier_std = init_xavier_std
+ # Hungarian matcher && loss
+ self.cross_entropy_weight = cross_entropy_weight
+ self.dice_weight = dice_weight
+ self.mask_weight = mask_weight
+ self.use_auxiliary_loss = use_auxiliary_loss
+ self.no_object_weight = no_object_weight
+ self.output_auxiliary_logits = output_auxiliary_logits
+
+ self.num_attention_heads = self.decoder_config.encoder_attention_heads
+ self.num_hidden_layers = self.decoder_config.num_hidden_layers
+ self.backbone = backbone
+ self.use_pretrained_backbone = use_pretrained_backbone
+ self.use_timm_backbone = use_timm_backbone
+ self.backbone_kwargs = backbone_kwargs
+ super().__init__(**kwargs)
+
+ @classmethod
+ def from_backbone_and_decoder_configs(
+ cls, backbone_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
+ ):
+ """Instantiate a [`MaskFormerConfig`] (or a derived class) from a pre-trained backbone model configuration and DETR model
+ configuration.
+
+ Args:
+ backbone_config ([`PretrainedConfig`]):
+ The backbone configuration.
+ decoder_config ([`PretrainedConfig`]):
+ The transformer decoder configuration to use.
+
+ Returns:
+ [`MaskFormerConfig`]: An instance of a configuration object
+ """
+ return cls(
+ backbone_config=backbone_config,
+ decoder_config=decoder_config,
+ **kwargs,
+ )
+
+
+__all__ = ["MaskFormerConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer_swin.py b/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3010e71d017d54e3dd782664c09fb4edb59fc4e
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/configuration_maskformer_swin.py
@@ -0,0 +1,153 @@
+# 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.
+"""MaskFormer Swin 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__)
+
+
+class MaskFormerSwinConfig(BackboneConfigMixin, PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MaskFormerSwinModel`]. 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 Swin
+ [microsoft/swin-tiny-patch4-window7-224](https://huggingface.co/microsoft/swin-tiny-patch4-window7-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:
+ 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-12):
+ The epsilon used by the layer normalization layers.
+ 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 MaskFormerSwinConfig, MaskFormerSwinModel
+
+ >>> # Initializing a microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> configuration = MaskFormerSwinConfig()
+
+ >>> # Initializing a model (with random weights) from the microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> model = MaskFormerSwinModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "maskformer-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,
+ 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.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))
+ 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
+ )
+
+
+__all__ = ["MaskFormerSwinConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/feature_extraction_maskformer.py b/lib/python3.12/site-packages/transformers/models/maskformer/feature_extraction_maskformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ce45471abd093581122ee4a269225c99e4d0c2b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/feature_extraction_maskformer.py
@@ -0,0 +1,36 @@
+# 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 MaskFormer."""
+
+import warnings
+
+from ...utils import logging
+from .image_processing_maskformer import MaskFormerImageProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+class MaskFormerFeatureExtractor(MaskFormerImageProcessor):
+ def __init__(self, *args, **kwargs) -> None:
+ warnings.warn(
+ "The class MaskFormerFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
+ " Please use MaskFormerImageProcessor instead.",
+ FutureWarning,
+ )
+ super().__init__(*args, **kwargs)
+
+
+__all__ = ["MaskFormerFeatureExtractor"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/image_processing_maskformer.py b/lib/python3.12/site-packages/transformers/models/maskformer/image_processing_maskformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..b31d0321881a453f66db72a1f152a0446f7d4726
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/image_processing_maskformer.py
@@ -0,0 +1,1278 @@
+# 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 MaskFormer."""
+
+import math
+import warnings
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import numpy as np
+
+from ...image_processing_utils import INIT_SERVICE_KWARGS, BaseImageProcessor, BatchFeature, get_size_dict
+from ...image_transforms import (
+ PaddingMode,
+ get_resize_output_image_size,
+ pad,
+ rescale,
+ resize,
+ to_channel_dimension_format,
+)
+from ...image_utils import (
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ get_image_size,
+ infer_channel_dimension_format,
+ is_scaled_image,
+ make_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import (
+ IMAGENET_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ TensorType,
+ filter_out_non_signature_kwargs,
+ is_torch_available,
+ is_torch_tensor,
+ logging,
+)
+from ...utils.deprecation import deprecate_kwarg
+
+
+logger = logging.get_logger(__name__)
+
+
+if TYPE_CHECKING:
+ from transformers import MaskFormerForInstanceSegmentationOutput
+
+
+if is_torch_available():
+ import torch
+ from torch import nn
+
+
+# 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.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
+
+
+# TODO: (Amy) Move to image_transforms
+def convert_segmentation_map_to_binary_masks(
+ segmentation_map: "np.ndarray",
+ instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
+ ignore_index: Optional[int] = None,
+ do_reduce_labels: bool = False,
+):
+ if do_reduce_labels and ignore_index is None:
+ raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.")
+
+ if do_reduce_labels:
+ segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1)
+
+ # Get unique ids (class or instance ids based on input)
+ all_labels = np.unique(segmentation_map)
+
+ # Drop background label if applicable
+ if ignore_index is not None:
+ all_labels = all_labels[all_labels != ignore_index]
+
+ # Generate a binary mask for each object instance
+ binary_masks = [(segmentation_map == i) for i in all_labels]
+
+ # Stack the binary masks
+ if binary_masks:
+ binary_masks = np.stack(binary_masks, axis=0)
+ else:
+ binary_masks = np.zeros((0, *segmentation_map.shape))
+
+ # Convert instance ids to class ids
+ if instance_id_to_semantic_id is not None:
+ labels = np.zeros(all_labels.shape[0])
+
+ for label in all_labels:
+ class_id = instance_id_to_semantic_id[label + 1 if do_reduce_labels else label]
+ labels[all_labels == label] = class_id - 1 if do_reduce_labels else class_id
+ else:
+ labels = all_labels
+
+ return binary_masks.astype(np.float32), labels.astype(np.int64)
+
+
+def get_maskformer_resize_output_image_size(
+ image: np.ndarray,
+ size: Union[int, Tuple[int, int], List[int], Tuple[int]],
+ max_size: Optional[int] = None,
+ size_divisor: int = 0,
+ default_to_square: bool = True,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+) -> Tuple[int, int]:
+ """
+ Computes the output size given the desired size.
+
+ Args:
+ image (`np.ndarray`):
+ The input image.
+ size (`int` or `Tuple[int, int]` or `List[int]` or `Tuple[int]`):
+ The size of the output image.
+ max_size (`int`, *optional*):
+ The maximum size of the output image.
+ size_divisor (`int`, *optional*, defaults to 0):
+ If `size_divisor` is given, the output image size will be divisible by the number.
+ default_to_square (`bool`, *optional*, defaults to `True`):
+ Whether to default to square if no size is provided.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format of the input image. If unset, will use the inferred format from the input.
+
+ Returns:
+ `Tuple[int, int]`: The output size.
+ """
+ output_size = get_resize_output_image_size(
+ input_image=image,
+ size=size,
+ default_to_square=default_to_square,
+ max_size=max_size,
+ input_data_format=input_data_format,
+ )
+
+ if size_divisor > 0:
+ height, width = output_size
+ height = int(math.ceil(height / size_divisor) * size_divisor)
+ width = int(math.ceil(width / size_divisor) * size_divisor)
+ output_size = (height, width)
+
+ return output_size
+
+
+class MaskFormerImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a MaskFormer image processor. The image processor can be used to prepare image(s) and optional targets
+ for the model.
+
+ This image processor inherits from [`BaseImageProcessor`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to resize the input to a certain `size`.
+ size (`int`, *optional*, defaults to 800):
+ Resize the input to the given size. Only has an effect if `do_resize` is set to `True`. If size is a
+ sequence like `(width, height)`, output size will be matched to this. If size is an int, smaller edge of
+ the image will be matched to this number. i.e, if `height > width`, then image will be rescaled to `(size *
+ height / width, size)`.
+ size_divisor (`int`, *optional*, defaults to 32):
+ Some backbones need images divisible by a certain number. If not passed, it defaults to the value used in
+ Swin Transformer.
+ resample (`int`, *optional*, defaults to `Resampling.BILINEAR`):
+ An optional resampling filter. This can be one of `PIL.Image.Resampling.NEAREST`,
+ `PIL.Image.Resampling.BOX`, `PIL.Image.Resampling.BILINEAR`, `PIL.Image.Resampling.HAMMING`,
+ `PIL.Image.Resampling.BICUBIC` or `PIL.Image.Resampling.LANCZOS`. Only has an effect if `do_resize` is set
+ to `True`.
+ do_rescale (`bool`, *optional*, defaults to `True`):
+ Whether to rescale the input to a certain `scale`.
+ rescale_factor (`float`, *optional*, defaults to `1/ 255`):
+ Rescale the input by the given factor. Only has an effect if `do_rescale` is set to `True`.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether or not to normalize the input with mean and standard deviation.
+ image_mean (`int`, *optional*, defaults to `[0.485, 0.456, 0.406]`):
+ The sequence of means for each channel, to be used when normalizing images. Defaults to the ImageNet mean.
+ image_std (`int`, *optional*, defaults to `[0.229, 0.224, 0.225]`):
+ The sequence of standard deviations for each channel, to be used when normalizing images. Defaults to the
+ ImageNet std.
+ ignore_index (`int`, *optional*):
+ Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels
+ denoted with 0 (background) will be replaced with `ignore_index`.
+ do_reduce_labels (`bool`, *optional*, defaults to `False`):
+ Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0
+ is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k).
+ The background label will be replaced by `ignore_index`.
+ num_labels (`int`, *optional*):
+ The number of labels in the segmentation map.
+
+ """
+
+ model_input_names = ["pixel_values", "pixel_mask"]
+
+ @deprecate_kwarg("reduce_labels", new_name="do_reduce_labels", version="4.44.0")
+ @deprecate_kwarg("size_divisibility", new_name="size_divisor", version="4.41.0")
+ @deprecate_kwarg("max_size", version="4.27.0", warn_if_greater_or_equal_version=True)
+ @filter_out_non_signature_kwargs(extra=["max_size", *INIT_SERVICE_KWARGS])
+ def __init__(
+ self,
+ do_resize: bool = True,
+ size: Dict[str, int] = None,
+ size_divisor: int = 32,
+ resample: PILImageResampling = PILImageResampling.BILINEAR,
+ do_rescale: bool = True,
+ rescale_factor: float = 1 / 255,
+ do_normalize: bool = True,
+ image_mean: Union[float, List[float]] = None,
+ image_std: Union[float, List[float]] = None,
+ ignore_index: Optional[int] = None,
+ do_reduce_labels: bool = False,
+ num_labels: Optional[int] = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ # We make max_size a private attribute so we can pass it as a default value in the preprocess method whilst
+ # `size` can still be pass in as an int
+ self._max_size = kwargs.pop("max_size", 1333)
+
+ size = size if size is not None else {"shortest_edge": 800, "longest_edge": self._max_size}
+ size = get_size_dict(size, max_size=self._max_size, default_to_square=False)
+
+ self.do_resize = do_resize
+ self.size = size
+ self.resample = resample
+ self.size_divisor = size_divisor
+ 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_DEFAULT_MEAN
+ self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
+ self.ignore_index = ignore_index
+ self.do_reduce_labels = do_reduce_labels
+ self.num_labels = num_labels
+
+ @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. `MaskFormerImageProcessor.from_pretrained(checkpoint, 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 "size_divisibility" in kwargs:
+ image_processor_dict["size_divisor"] = kwargs.pop("size_divisibility")
+ if "reduce_labels" in image_processor_dict:
+ image_processor_dict["do_reduce_labels"] = image_processor_dict.pop("reduce_labels")
+ return super().from_dict(image_processor_dict, **kwargs)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the
+ `_max_size` attribute from the dictionary.
+ """
+ image_processor_dict = super().to_dict()
+ image_processor_dict.pop("_max_size", None)
+ return image_processor_dict
+
+ @deprecate_kwarg("max_size", version="4.27.0", warn_if_greater_or_equal_version=True)
+ def resize(
+ self,
+ image: np.ndarray,
+ size: Dict[str, int],
+ size_divisor: int = 0,
+ resample: PILImageResampling = PILImageResampling.BILINEAR,
+ data_format=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 size of the output image.
+ size_divisor (`int`, *optional*, defaults to 0):
+ If `size_divisor` is given, the output image size will be divisible by the number.
+ resample (`PILImageResampling` resampling filter, *optional*, defaults to `PILImageResampling.BILINEAR`):
+ Resampling filter to use when resizing the image.
+ 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.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format of the input image. If not provided, it will be inferred.
+ """
+
+ # Deprecated, backward compatibility
+ max_size = kwargs.pop("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, max_size = size["shortest_edge"], size["longest_edge"]
+ elif "height" in size and "width" in size:
+ size = (size["height"], size["width"])
+ max_size = None
+ else:
+ raise ValueError(
+ "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
+ f" {size.keys()}."
+ )
+ size = get_maskformer_resize_output_image_size(
+ image=image,
+ size=size,
+ max_size=max_size,
+ size_divisor=size_divisor,
+ default_to_square=False,
+ input_data_format=input_data_format,
+ )
+ 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.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)
+
+ def convert_segmentation_map_to_binary_masks(
+ self,
+ segmentation_map: "np.ndarray",
+ instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
+ ignore_index: Optional[int] = None,
+ do_reduce_labels: bool = False,
+ ):
+ do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels
+ ignore_index = ignore_index if ignore_index is not None else self.ignore_index
+ return convert_segmentation_map_to_binary_masks(
+ segmentation_map=segmentation_map,
+ instance_id_to_semantic_id=instance_id_to_semantic_id,
+ ignore_index=ignore_index,
+ do_reduce_labels=do_reduce_labels,
+ )
+
+ def __call__(self, images, segmentation_maps=None, **kwargs) -> BatchFeature:
+ return self.preprocess(images, segmentation_maps=segmentation_maps, **kwargs)
+
+ def _preprocess(
+ self,
+ image: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ size_divisor: Optional[int] = None,
+ resample: PILImageResampling = 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,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ):
+ if do_resize:
+ image = self.resize(
+ image, size=size, size_divisor=size_divisor, resample=resample, input_data_format=input_data_format
+ )
+ if do_rescale:
+ image = self.rescale(image, rescale_factor=rescale_factor, input_data_format=input_data_format)
+ if do_normalize:
+ image = self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+ return image
+
+ def _preprocess_image(
+ self,
+ image: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ size_divisor: Optional[int] = None,
+ resample: PILImageResampling = 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,
+ data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> np.ndarray:
+ """Preprocesses a single image."""
+ # All transformations expect numpy arrays.
+ image = to_numpy_array(image)
+ if do_rescale and is_scaled_image(image):
+ 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:
+ input_data_format = infer_channel_dimension_format(image)
+ image = self._preprocess(
+ image=image,
+ do_resize=do_resize,
+ size=size,
+ size_divisor=size_divisor,
+ resample=resample,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ input_data_format=input_data_format,
+ )
+ if data_format is not None:
+ image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+ return image
+
+ def _preprocess_mask(
+ self,
+ segmentation_map: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ size_divisor: int = 0,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> np.ndarray:
+ """Preprocesses a single mask."""
+ segmentation_map = to_numpy_array(segmentation_map)
+ # Add channel dimension if missing - needed for certain transformations
+ if segmentation_map.ndim == 2:
+ added_channel_dim = True
+ segmentation_map = segmentation_map[None, ...]
+ input_data_format = ChannelDimension.FIRST
+ else:
+ added_channel_dim = False
+ if input_data_format is None:
+ input_data_format = infer_channel_dimension_format(segmentation_map, num_channels=1)
+ # TODO: (Amy)
+ # Remork segmentation map processing to include reducing labels and resizing which doesn't
+ # drop segment IDs > 255.
+ segmentation_map = self._preprocess(
+ image=segmentation_map,
+ do_resize=do_resize,
+ resample=PILImageResampling.NEAREST,
+ size=size,
+ size_divisor=size_divisor,
+ do_rescale=False,
+ do_normalize=False,
+ input_data_format=input_data_format,
+ )
+ # Remove extra channel dimension if added for processing
+ if added_channel_dim:
+ segmentation_map = segmentation_map.squeeze(0)
+ return segmentation_map
+
+ @deprecate_kwarg("reduce_labels", new_name="do_reduce_labels", version="4.44.0")
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ segmentation_maps: Optional[ImageInput] = None,
+ instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
+ do_resize: Optional[bool] = None,
+ size: Optional[Dict[str, int]] = None,
+ size_divisor: Optional[int] = None,
+ resample: PILImageResampling = 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,
+ ignore_index: Optional[int] = None,
+ do_reduce_labels: Optional[bool] = None,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> BatchFeature:
+ 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, default_to_square=False, max_size=self._max_size)
+ size_divisor = size_divisor if size_divisor is not None else self.size_divisor
+ 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
+ ignore_index = ignore_index if ignore_index is not None else self.ignore_index
+ do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels
+
+ 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,
+ )
+
+ if segmentation_maps is not None and not valid_images(segmentation_maps):
+ raise ValueError(
+ "Invalid segmentation map type. Must be of type PIL.Image.Image, numpy.ndarray, "
+ "torch.Tensor, tf.Tensor or jax.ndarray."
+ )
+
+ images = make_list_of_images(images)
+ if segmentation_maps is not None:
+ segmentation_maps = make_list_of_images(segmentation_maps, expected_ndims=2)
+
+ if segmentation_maps is not None and len(images) != len(segmentation_maps):
+ raise ValueError("Images and segmentation maps must have the same length.")
+
+ images = [
+ self._preprocess_image(
+ image,
+ do_resize=do_resize,
+ size=size,
+ size_divisor=size_divisor,
+ resample=resample,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+ for image in images
+ ]
+
+ if segmentation_maps is not None:
+ segmentation_maps = [
+ self._preprocess_mask(
+ segmentation_map, do_resize, size, size_divisor, input_data_format=input_data_format
+ )
+ for segmentation_map in segmentation_maps
+ ]
+ encoded_inputs = self.encode_inputs(
+ images,
+ segmentation_maps,
+ instance_id_to_semantic_id,
+ ignore_index,
+ do_reduce_labels,
+ return_tensors,
+ input_data_format=data_format,
+ )
+ return encoded_inputs
+
+ # Copied from transformers.models.vilt.image_processing_vilt.ViltImageProcessor._pad_image
+ def _pad_image(
+ self,
+ image: np.ndarray,
+ output_size: Tuple[int, int],
+ constant_values: Union[float, Iterable[float]] = 0,
+ data_format: Optional[ChannelDimension] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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,
+ )
+ return padded_image
+
+ # Copied from transformers.models.vilt.image_processing_vilt.ViltImageProcessor.pad
+ def pad(
+ self,
+ images: List[np.ndarray],
+ 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,
+ ) -> 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:
+ image (`np.ndarray`):
+ Image to pad.
+ 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.
+ """
+ pad_size = get_max_height_width(images, input_data_format=input_data_format)
+
+ padded_images = [
+ self._pad_image(
+ image,
+ pad_size,
+ constant_values=constant_values,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+ for image in images
+ ]
+ 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
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+ def encode_inputs(
+ self,
+ pixel_values_list: List[ImageInput],
+ segmentation_maps: ImageInput = None,
+ instance_id_to_semantic_id: Optional[Union[List[Dict[int, int]], Dict[int, int]]] = None,
+ ignore_index: Optional[int] = None,
+ do_reduce_labels: bool = False,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ):
+ """
+ Pad images up to the largest image in a batch and create a corresponding `pixel_mask`.
+
+ MaskFormer addresses semantic segmentation with a mask classification paradigm, thus input segmentation maps
+ will be converted to lists of binary masks and their respective labels. Let's see an example, assuming
+ `segmentation_maps = [[2,6,7,9]]`, the output will contain `mask_labels =
+ [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]]` (four binary masks) and `class_labels = [2,6,7,9]`, the labels for
+ each mask.
+
+ Args:
+ pixel_values_list (`List[ImageInput]`):
+ List of images (pixel values) to be padded. Each image should be a tensor of shape `(channels, height,
+ width)`.
+
+ segmentation_maps (`ImageInput`, *optional*):
+ The corresponding semantic segmentation maps with the pixel-wise annotations.
+
+ (`bool`, *optional*, defaults to `True`):
+ Whether or not to pad images up to the largest image in a batch and create a pixel mask.
+
+ If left to the default, will return a pixel mask that is:
+
+ - 1 for pixels that are real (i.e. **not masked**),
+ - 0 for pixels that are padding (i.e. **masked**).
+
+ instance_id_to_semantic_id (`List[Dict[int, int]]` or `Dict[int, int]`, *optional*):
+ A mapping between object instance ids and class ids. If passed, `segmentation_maps` is treated as an
+ instance segmentation map where each pixel represents an instance id. Can be provided as a single
+ dictionary with a global/dataset-level mapping or as a list of dictionaries (one per image), to map
+ instance ids in each image separately.
+
+ return_tensors (`str` or [`~file_utils.TensorType`], *optional*):
+ If set, will return tensors instead of NumPy arrays. If set to `'pt'`, return PyTorch `torch.Tensor`
+ objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **pixel_values** -- Pixel values to be fed to a model.
+ - **pixel_mask** -- Pixel mask to be fed to a model (when `=True` or if `pixel_mask` is in
+ `self.model_input_names`).
+ - **mask_labels** -- Optional list of mask labels of shape `(labels, height, width)` to be fed to a model
+ (when `annotations` are provided).
+ - **class_labels** -- Optional list of class labels of shape `(labels)` to be fed to a model (when
+ `annotations` are provided). They identify the labels of `mask_labels`, e.g. the label of
+ `mask_labels[i][j]` if `class_labels[i][j]`.
+ """
+ ignore_index = self.ignore_index if ignore_index is None else ignore_index
+ do_reduce_labels = self.do_reduce_labels if do_reduce_labels is None else do_reduce_labels
+
+ pixel_values_list = [to_numpy_array(pixel_values) for pixel_values in pixel_values_list]
+
+ if input_data_format is None:
+ input_data_format = infer_channel_dimension_format(pixel_values_list[0])
+
+ encoded_inputs = self.pad(
+ pixel_values_list, return_tensors=return_tensors, input_data_format=input_data_format
+ )
+
+ if segmentation_maps is not None:
+ mask_labels = []
+ class_labels = []
+ pad_size = get_max_height_width(pixel_values_list, input_data_format=input_data_format)
+ # Convert to list of binary masks and labels
+ for idx, segmentation_map in enumerate(segmentation_maps):
+ segmentation_map = to_numpy_array(segmentation_map)
+ if isinstance(instance_id_to_semantic_id, list):
+ instance_id = instance_id_to_semantic_id[idx]
+ else:
+ instance_id = instance_id_to_semantic_id
+ # Use instance2class_id mapping per image
+ masks, classes = self.convert_segmentation_map_to_binary_masks(
+ segmentation_map, instance_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels
+ )
+ # We add an axis to make them compatible with the transformations library
+ # this will be removed in the future
+ if masks.shape[0] > 0:
+ masks = [mask[None, ...] for mask in masks]
+ masks = [
+ self._pad_image(
+ image=mask,
+ output_size=pad_size,
+ constant_values=ignore_index,
+ input_data_format=ChannelDimension.FIRST,
+ )
+ for mask in masks
+ ]
+ masks = np.concatenate(masks, axis=0)
+ else:
+ masks = np.zeros((0, *pad_size), dtype=np.float32)
+ mask_labels.append(torch.from_numpy(masks))
+ class_labels.append(torch.from_numpy(classes))
+
+ # we cannot batch them since they don't share a common class size
+ encoded_inputs["mask_labels"] = mask_labels
+ encoded_inputs["class_labels"] = class_labels
+
+ return encoded_inputs
+
+ def post_process_segmentation(
+ self, outputs: "MaskFormerForInstanceSegmentationOutput", target_size: Tuple[int, int] = None
+ ) -> "torch.Tensor":
+ """
+ Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image segmentation predictions. Only
+ supports PyTorch.
+
+ Args:
+ outputs ([`MaskFormerForInstanceSegmentationOutput`]):
+ The outputs from [`MaskFormerForInstanceSegmentation`].
+
+ target_size (`Tuple[int, int]`, *optional*):
+ If set, the `masks_queries_logits` will be resized to `target_size`.
+
+ Returns:
+ `torch.Tensor`:
+ A tensor of shape (`batch_size, num_class_labels, height, width`).
+ """
+ warnings.warn(
+ "`post_process_segmentation` is deprecated and will be removed in v5 of Transformers, please use"
+ " `post_process_instance_segmentation`",
+ FutureWarning,
+ )
+
+ # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
+ class_queries_logits = outputs.class_queries_logits
+ # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
+ masks_queries_logits = outputs.masks_queries_logits
+ if target_size is not None:
+ masks_queries_logits = torch.nn.functional.interpolate(
+ masks_queries_logits,
+ size=target_size,
+ mode="bilinear",
+ align_corners=False,
+ )
+ # remove the null class `[..., :-1]`
+ masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
+ # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
+ masks_probs = masks_queries_logits.sigmoid()
+ # now we want to sum over the queries,
+ # $ out_{c,h,w} = \sum_q p_{q,c} * m_{q,h,w} $
+ # where $ softmax(p) \in R^{q, c} $ is the mask classes
+ # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
+ # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
+ segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
+
+ return segmentation
+
+ def post_process_semantic_segmentation(
+ self, outputs, target_sizes: Optional[List[Tuple[int, int]]] = None
+ ) -> "torch.Tensor":
+ """
+ Converts the output of [`MaskFormerForInstanceSegmentation`] into semantic segmentation maps. Only supports
+ PyTorch.
+
+ Args:
+ outputs ([`MaskFormerForInstanceSegmentation`]):
+ Raw outputs of the model.
+ target_sizes (`List[Tuple[int, int]]`, *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 left to None, 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.class_queries_logits # [batch_size, num_queries, num_classes+1]
+ masks_queries_logits = outputs.masks_queries_logits # [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 = torch.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
+
+ 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,
+ return_binary_maps: Optional[bool] = False,
+ ) -> List[Dict]:
+ """
+ Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into instance segmentation predictions. Only
+ supports PyTorch. If instances could overlap, set either return_coco_annotation or return_binary_maps
+ to `True` to get the correct segmentation result.
+
+ Args:
+ outputs ([`MaskFormerForInstanceSegmentation`]):
+ 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 left to None, 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.
+ return_binary_maps (`bool`, *optional*, defaults to `False`):
+ If set to `True`, segmentation maps are returned as a concatenated tensor of binary segmentation maps
+ (one per detected instance).
+ 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`, or a tensor of shape `(num_instances, height, width)` if return_binary_maps 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`.
+ """
+ if return_coco_annotation and return_binary_maps:
+ raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.")
+
+ # [batch_size, num_queries, num_classes+1]
+ class_queries_logits = outputs.class_queries_logits
+ # [batch_size, num_queries, height, width]
+ masks_queries_logits = outputs.masks_queries_logits
+
+ device = masks_queries_logits.device
+ num_classes = class_queries_logits.shape[-1] - 1
+ num_queries = class_queries_logits.shape[-2]
+
+ # Loop over items in batch size
+ results: List[Dict[str, TensorType]] = []
+
+ for i in range(class_queries_logits.shape[0]):
+ mask_pred = masks_queries_logits[i]
+ mask_cls = class_queries_logits[i]
+
+ scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1]
+ labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
+
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False)
+ labels_per_image = labels[topk_indices]
+
+ topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor")
+ mask_pred = mask_pred[topk_indices]
+ pred_masks = (mask_pred > 0).float()
+
+ # Calculate average mask prob
+ mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / (
+ pred_masks.flatten(1).sum(1) + 1e-6
+ )
+ pred_scores = scores_per_image * mask_scores_per_image
+ pred_classes = labels_per_image
+
+ segmentation = torch.zeros(masks_queries_logits.shape[2:]) - 1
+ if target_sizes is not None:
+ segmentation = torch.zeros(target_sizes[i]) - 1
+ pred_masks = torch.nn.functional.interpolate(
+ pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest"
+ )[0]
+
+ instance_maps, segments = [], []
+ current_segment_id = 0
+ for j in range(num_queries):
+ score = pred_scores[j].item()
+
+ if not torch.all(pred_masks[j] == 0) and score >= threshold:
+ segmentation[pred_masks[j] == 1] = current_segment_id
+ segments.append(
+ {
+ "id": current_segment_id,
+ "label_id": pred_classes[j].item(),
+ "was_fused": False,
+ "score": round(score, 6),
+ }
+ )
+ current_segment_id += 1
+ instance_maps.append(pred_masks[j])
+
+ # Return segmentation map in run-length encoding (RLE) format
+ if return_coco_annotation:
+ segmentation = convert_segmentation_to_rle(segmentation)
+
+ # Return a concatenated tensor of binary instance maps
+ if return_binary_maps and len(instance_maps) != 0:
+ segmentation = torch.stack(instance_maps, dim=0)
+
+ results.append({"segmentation": segmentation, "segments_info": segments})
+ return results
+
+ 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 [`MaskFormerForInstanceSegmentationOutput`] into image panoptic segmentation
+ predictions. Only supports PyTorch.
+
+ Args:
+ outputs ([`MaskFormerForInstanceSegmentationOutput`]):
+ The outputs from [`MaskFormerForInstanceSegmentation`].
+ 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 left to None, 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`, set
+ to `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("`label_ids_to_fuse` unset. No instance will be fused.")
+ label_ids_to_fuse = set()
+
+ class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
+ masks_queries_logits = outputs.masks_queries_logits # [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
+
+
+__all__ = ["MaskFormerImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer.py b/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c1873b4d6663209277ad56d147d1bf835548a8c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer.py
@@ -0,0 +1,1885 @@
+# coding=utf-8
+# Copyright 2022 Meta Platforms, Inc.s 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 MaskFormer model."""
+
+import math
+from dataclasses import dataclass
+from numbers import Number
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+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 BaseModelOutputWithCrossAttentions
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import is_torch_greater_or_equal_than_2_1
+from ...utils import (
+ ModelOutput,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_accelerate_available,
+ is_scipy_available,
+ logging,
+ replace_return_docstrings,
+ requires_backends,
+)
+from ...utils.backbone_utils import load_backbone
+from ...utils.import_utils import is_torchdynamo_compiling
+from ..detr import DetrConfig
+from .configuration_maskformer import MaskFormerConfig
+from .configuration_maskformer_swin import MaskFormerSwinConfig
+
+
+if is_accelerate_available():
+ from accelerate import PartialState
+ from accelerate.utils import reduce
+
+if is_scipy_available():
+ from scipy.optimize import linear_sum_assignment
+
+logger = logging.get_logger(__name__)
+
+
+_CONFIG_FOR_DOC = "MaskFormerConfig"
+_CHECKPOINT_FOR_DOC = "facebook/maskformer-swin-base-ade"
+
+
+@dataclass
+# Copied from transformers.models.detr.modeling_detr.DetrDecoderOutput
+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 MaskFormerPixelLevelModuleOutput(ModelOutput):
+ """
+ MaskFormer's pixel level module output. It returns both the last and (optionally) the hidden states from the
+ `encoder` and `decoder`. By default, the `encoder` is a MaskFormerSwin Transformer and the `decoder` is a Feature
+ Pyramid Network (FPN).
+
+ The `encoder_last_hidden_state` are referred on the paper as **images features**, while `decoder_last_hidden_state`
+ as **pixel embeddings**
+
+ Args:
+ encoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage of the encoder.
+ 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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the model at
+ the output of each stage.
+ decoder_last_hidden_state (`torch.FloatTensor` of shape`(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage 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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the model at
+ the output of each stage.
+ """
+
+ encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ decoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class MaskFormerPixelDecoderOutput(ModelOutput):
+ """
+ MaskFormer's pixel decoder module output, practically a Feature Pyramid Network. It returns the last hidden state
+ and (optionally) the hidden states.
+
+ Args:
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage 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, num_channels, height, width)`. 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 from Detr's decoder after the attention softmax, used to compute the
+ weighted average in the self-attention heads.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class MaskFormerModelOutput(ModelOutput):
+ """
+ Class for outputs of [`MaskFormerModel`]. This class returns all the needed hidden states to compute the logits.
+
+ Args:
+ encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage of the encoder model (backbone).
+ pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN).
+ transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Last hidden states (final feature map) of the last stage of the transformer decoder 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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder
+ model at the output of each stage.
+ pixel_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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel
+ decoder model at the output of each stage.
+ transformer_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 stage) of
+ shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the
+ transformer decoder at the output of each stage.
+ 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and
+ `decoder_hidden_states`
+ 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 from Detr's decoder after the attention softmax, used to compute the
+ weighted average in the self-attention heads.
+ """
+
+ encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ pixel_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ transformer_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ pixel_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ transformer_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class MaskFormerForInstanceSegmentationOutput(ModelOutput):
+ """
+ Class for outputs of [`MaskFormerForInstanceSegmentation`].
+
+ This output can be directly passed to [`~MaskFormerImageProcessor.post_process_semantic_segmentation`] or or
+ [`~MaskFormerImageProcessor.post_process_instance_segmentation`] or
+ [`~MaskFormerImageProcessor.post_process_panoptic_segmentation`] depending on the task. Please, see
+ [`~MaskFormerImageProcessor] for details regarding usage.
+
+ Args:
+ loss (`torch.Tensor`, *optional*):
+ The computed loss, returned when labels are present.
+ class_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each
+ query. Note the `+ 1` is needed because we incorporate the null class.
+ masks_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each
+ query.
+ encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage of the encoder model (backbone).
+ pixel_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last stage of the pixel decoder model (FPN).
+ transformer_decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Last hidden states (final feature map) of the last stage of the transformer decoder 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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the encoder
+ model at the output of each stage.
+ pixel_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 stage) of
+ shape `(batch_size, num_channels, height, width)`. Hidden-states (also called feature maps) of the pixel
+ decoder model at the output of each stage.
+ transformer_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 stage) of
+ shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the transformer decoder at the output
+ of each stage.
+ 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` containing `encoder_hidden_states`, `pixel_decoder_hidden_states` and
+ `decoder_hidden_states`.
+ 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 from Detr's decoder after the attention softmax, used to compute the
+ weighted average in the self-attention heads.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ class_queries_logits: Optional[torch.FloatTensor] = None
+ masks_queries_logits: Optional[torch.FloatTensor] = None
+ auxiliary_logits: Optional[torch.FloatTensor] = None
+ encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ pixel_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ transformer_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ pixel_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ transformer_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+def upsample_like(pixel_values: Tensor, like: Tensor, mode: str = "bilinear") -> Tensor:
+ """
+ An utility function that upsamples `pixel_values` to match the dimension of `like`.
+
+ Args:
+ pixel_values (`torch.Tensor`):
+ The tensor we wish to upsample.
+ like (`torch.Tensor`):
+ The tensor we wish to use as size target.
+ mode (str, *optional*, defaults to `"bilinear"`):
+ The interpolation mode.
+
+ Returns:
+ `torch.Tensor`: The upsampled tensor
+ """
+ _, _, height, width = like.shape
+ upsampled = nn.functional.interpolate(pixel_values, size=(height, width), mode=mode, align_corners=False)
+ return upsampled
+
+
+# refactored from original implementation
+def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor:
+ r"""
+ Compute the DICE loss, similar to generalized IOU for masks as follows:
+
+ $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$
+
+ In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow
+
+ $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ num_masks (`int`):
+ The number of masks present in the current batch, used for normalization.
+
+ Returns:
+ `torch.Tensor`: The computed loss.
+ """
+ probs = inputs.sigmoid().flatten(1)
+ numerator = 2 * (probs * labels).sum(-1)
+ denominator = probs.sum(-1) + labels.sum(-1)
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ loss = loss.sum() / num_masks
+ return loss
+
+
+# refactored from original implementation
+def sigmoid_focal_loss(
+ inputs: Tensor, labels: Tensor, num_masks: int, alpha: float = 0.25, gamma: float = 2
+) -> Tensor:
+ r"""
+ Focal loss proposed in [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002) originally used in
+ RetinaNet. The loss is computed as follows:
+
+ $$ \mathcal{L}_{\text{focal loss} = -(1 - p_t)^{\gamma}\log{(p_t)} $$
+
+ where \\(CE(p_t) = -\log{(p_t)}}\\), CE is the standard Cross Entropy Loss
+
+ Please refer to equation (1,2,3) of the paper for a better understanding.
+
+ Args:
+ inputs (`torch.Tensor`):
+ A float tensor of arbitrary shape.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ num_masks (`int`):
+ The number of masks present in the current batch, used for normalization.
+ alpha (float, *optional*, defaults to 0.25):
+ Weighting factor in range (0,1) to balance positive vs negative examples.
+ gamma (float, *optional*, defaults to 2.0):
+ Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples.
+
+ Returns:
+ `torch.Tensor`: The computed loss.
+ """
+ criterion = nn.BCEWithLogitsLoss(reduction="none")
+ probs = inputs.sigmoid()
+ cross_entropy_loss = criterion(inputs, labels)
+ p_t = probs * labels + (1 - probs) * (1 - labels)
+ loss = cross_entropy_loss * ((1 - p_t) ** gamma)
+
+ if alpha >= 0:
+ alpha_t = alpha * labels + (1 - alpha) * (1 - labels)
+ loss = alpha_t * loss
+
+ loss = loss.mean(1).sum() / num_masks
+ return loss
+
+
+# refactored from original implementation
+def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor:
+ """
+ A pair wise version of the dice loss, see `dice_loss` for usage.
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+
+ Returns:
+ `torch.Tensor`: The computed loss between each pairs.
+ """
+ inputs = inputs.sigmoid().flatten(1)
+ numerator = 2 * torch.matmul(inputs, labels.T)
+ # using broadcasting to get a [num_queries, NUM_CLASSES] matrix
+ denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :]
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss
+
+
+# refactored from original implementation
+def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float = 0.25, gamma: float = 2.0) -> Tensor:
+ r"""
+ A pair wise version of the focal loss, see `sigmoid_focal_loss` for usage.
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ alpha (float, *optional*, defaults to 0.25):
+ Weighting factor in range (0,1) to balance positive vs negative examples.
+ gamma (float, *optional*, defaults to 2.0):
+ Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples.
+
+ Returns:
+ `torch.Tensor`: The computed loss between each pairs.
+ """
+ if alpha < 0:
+ raise ValueError("alpha must be positive")
+
+ height_and_width = inputs.shape[1]
+
+ criterion = nn.BCEWithLogitsLoss(reduction="none")
+ prob = inputs.sigmoid()
+ cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs))
+ focal_pos = ((1 - prob) ** gamma) * cross_entropy_loss_pos
+ focal_pos *= alpha
+
+ cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs))
+
+ focal_neg = (prob**gamma) * cross_entropy_loss_neg
+ focal_neg *= 1 - alpha
+
+ loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T)
+
+ return loss / height_and_width
+
+
+# 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]):
+ 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,
+ ) -> 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
+ 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
+
+
+# Copied from transformers.models.detr.modeling_detr.DetrDecoderLayer
+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,
+ ):
+ """
+ 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.
+ """
+ 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 DetrDecoder(nn.Module):
+ """
+ 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__()
+ self.config = 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
+
+ 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,
+ ):
+ 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.
+ """
+ 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:
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_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,
+ None,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ None,
+ output_attentions,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=None,
+ 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,
+ )
+
+
+# refactored from original implementation
+class MaskFormerHungarianMatcher(nn.Module):
+ """This class computes an assignment between the labels and the predictions of the network.
+
+ For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more
+ predictions than labels. 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).
+ """
+
+ def __init__(self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0):
+ """Creates the matcher
+
+ Params:
+ cost_class (float, *optional*, defaults to 1.0):
+ This is the relative weight of the classification error in the matching cost.
+ cost_mask (float, *optional*, defaults to 1.0):
+ This is the relative weight of the focal loss of the binary mask in the matching cost.
+ cost_dice (float, *optional*, defaults to 1.0):
+ This is the relative weight of the dice loss of the binary mask in the matching cost
+ """
+ super().__init__()
+ if cost_class == 0 and cost_mask == 0 and cost_dice == 0:
+ raise ValueError("All costs cant be 0")
+ self.cost_class = cost_class
+ self.cost_mask = cost_mask
+ self.cost_dice = cost_dice
+
+ @torch.no_grad()
+ def forward(self, masks_queries_logits, class_queries_logits, mask_labels, class_labels) -> List[Tuple[Tensor]]:
+ """Performs the matching
+
+ Params:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor` of dim `batch_size, num_queries, num_labels` with the
+ classification logits.
+ class_queries_logits (`torch.Tensor`):
+ A tensor` of dim `batch_size, num_queries, height, width` with the
+ predicted masks.
+
+ class_labels (`torch.Tensor`):
+ A tensor` of dim `num_target_boxes` (where num_target_boxes is the number
+ of ground-truth objects in the target) containing the class labels.
+ mask_labels (`torch.Tensor`):
+ A tensor` of dim `num_target_boxes, height, width` containing the target
+ masks.
+
+ Returns:
+ `List[Tuple[Tensor]]`: 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 labels (in order)
+ For each batch element, it holds:
+ len(index_i) = len(index_j) = min(num_queries, num_target_boxes).
+ """
+ indices: List[Tuple[np.array]] = []
+
+ preds_masks = masks_queries_logits
+ preds_probs = class_queries_logits
+ # iterate through batch size
+ for pred_probs, pred_mask, target_mask, labels in zip(preds_probs, preds_masks, mask_labels, class_labels):
+ # downsample the target mask, save memory
+ target_mask = nn.functional.interpolate(target_mask[:, None], size=pred_mask.shape[-2:], mode="nearest")
+ pred_probs = pred_probs.softmax(-1)
+ # 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.
+ cost_class = -pred_probs[:, labels]
+ # flatten spatial dimension "q h w -> q (h w)"
+ pred_mask_flat = pred_mask.flatten(1) # [num_queries, height*width]
+ # same for target_mask "c h w -> c (h w)"
+ target_mask_flat = target_mask[:, 0].flatten(1) # [num_total_labels, height*width]
+ # compute the focal loss between each mask pairs -> shape (num_queries, num_labels)
+ cost_mask = pair_wise_sigmoid_focal_loss(pred_mask_flat, target_mask_flat)
+ # Compute the dice loss betwen each mask pairs -> shape (num_queries, num_labels)
+ cost_dice = pair_wise_dice_loss(pred_mask_flat, target_mask_flat)
+ # final cost matrix
+ cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
+ # do the assigmented using the hungarian algorithm in scipy
+ assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
+ indices.append(assigned_indices)
+
+ # It could be stacked in one tensor
+ matched_indices = [
+ (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices
+ ]
+ return matched_indices
+
+ def __repr__(self):
+ head = "Matcher " + self.__class__.__name__
+ body = [
+ f"cost_class: {self.cost_class}",
+ f"cost_mask: {self.cost_mask}",
+ f"cost_dice: {self.cost_dice}",
+ ]
+ _repr_indent = 4
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)
+
+
+# copied and adapted from original implementation
+class MaskFormerLoss(nn.Module):
+ def __init__(
+ self,
+ num_labels: int,
+ matcher: MaskFormerHungarianMatcher,
+ weight_dict: Dict[str, float],
+ eos_coef: float,
+ ):
+ """
+ The MaskFormer Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we compute
+ hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair of
+ matched ground-truth / prediction (supervise class and mask)
+
+ Args:
+ num_labels (`int`):
+ The number of classes.
+ matcher (`MaskFormerHungarianMatcher`):
+ A torch module that computes the assigments between the predictions and labels.
+ weight_dict (`Dict[str, float]`):
+ A dictionary of weights to be applied to the different losses.
+ eos_coef (`float`):
+ Weight to apply to the null class.
+ """
+
+ super().__init__()
+ requires_backends(self, ["scipy"])
+ self.num_labels = num_labels
+ self.matcher = matcher
+ self.weight_dict = weight_dict
+ self.eos_coef = eos_coef
+ empty_weight = torch.ones(self.num_labels + 1)
+ empty_weight[-1] = self.eos_coef
+ self.register_buffer("empty_weight", empty_weight)
+
+ def _max_by_axis(self, the_list: 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
+
+ def _pad_images_to_max_in_batch(self, tensors: List[Tensor]) -> Tuple[Tensor, Tensor]:
+ # get the maximum size in the batch
+ max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors])
+ batch_size = len(tensors)
+ # compute finel size
+ batch_shape = [batch_size] + max_size
+ b, _, h, w = batch_shape
+ # get metadata
+ dtype = tensors[0].dtype
+ device = tensors[0].device
+ padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device)
+ padding_masks = torch.ones((b, h, w), dtype=torch.bool, device=device)
+ # pad the tensors to the size of the biggest one
+ for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks):
+ padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor)
+ padding_mask[: tensor.shape[1], : tensor.shape[2]] = False
+
+ return padded_tensors, padding_masks
+
+ def loss_labels(
+ self, class_queries_logits: Tensor, class_labels: List[Tensor], indices: Tuple[np.array]
+ ) -> Dict[str, Tensor]:
+ """Compute the losses related to the labels using cross entropy.
+
+ Args:
+ class_queries_logits (`torch.Tensor`):
+ A tensor of shape `batch_size, num_queries, num_labels`
+ class_labels (`List[torch.Tensor]`):
+ List of class labels of shape `(labels)`.
+ indices (`Tuple[np.array])`:
+ The indices computed by the Hungarian matcher.
+
+ Returns:
+ `Dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key:
+ - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels.
+ """
+
+ pred_logits = class_queries_logits
+ batch_size, num_queries, _ = pred_logits.shape
+ criterion = nn.CrossEntropyLoss(weight=self.empty_weight)
+ idx = self._get_predictions_permutation_indices(indices)
+ # shape = (batch_size, num_queries)
+ target_classes_o = torch.cat([target[j] for target, (_, j) in zip(class_labels, indices)])
+ # shape = (batch_size, num_queries)
+ target_classes = torch.full(
+ (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device
+ )
+ target_classes[idx] = target_classes_o
+ # target_classes is a (batch_size, num_labels, num_queries), we need to permute pred_logits "b q c -> b c q"
+ pred_logits_transposed = pred_logits.transpose(1, 2)
+ loss_ce = criterion(pred_logits_transposed, target_classes)
+ losses = {"loss_cross_entropy": loss_ce}
+ return losses
+
+ def loss_masks(
+ self, masks_queries_logits: Tensor, mask_labels: List[Tensor], indices: Tuple[np.array], num_masks: int
+ ) -> Dict[str, Tensor]:
+ """Compute the losses related to the masks using focal and dice loss.
+
+ Args:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor of shape `batch_size, num_queries, height, width`
+ mask_labels (`torch.Tensor`):
+ List of mask labels of shape `(labels, height, width)`.
+ indices (`Tuple[np.array])`:
+ The indices computed by the Hungarian matcher.
+ num_masks (`int)`:
+ The number of masks, used for normalization.
+
+ Returns:
+ `Dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys:
+ - **loss_mask** -- The loss computed using sigmoid focal loss on the predicted and ground truth masks.
+ - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth
+ masks.
+ """
+ src_idx = self._get_predictions_permutation_indices(indices)
+ tgt_idx = self._get_targets_permutation_indices(indices)
+ # shape (batch_size * num_queries, height, width)
+ pred_masks = masks_queries_logits[src_idx]
+ # shape (batch_size, num_queries, height, width)
+ # pad all and stack the targets to the num_labels dimension
+ target_masks, _ = self._pad_images_to_max_in_batch(mask_labels)
+ target_masks = target_masks[tgt_idx]
+ # upsample predictions to the target size, we have to add one dim to use interpolate
+ pred_masks = nn.functional.interpolate(
+ pred_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
+ )
+ pred_masks = pred_masks[:, 0].flatten(1)
+
+ target_masks = target_masks.flatten(1)
+ losses = {
+ "loss_mask": sigmoid_focal_loss(pred_masks, target_masks, num_masks),
+ "loss_dice": dice_loss(pred_masks, target_masks, num_masks),
+ }
+ return losses
+
+ def _get_predictions_permutation_indices(self, indices):
+ # permute predictions following indices
+ batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+ predictions_indices = torch.cat([src for (src, _) in indices])
+ return batch_indices, predictions_indices
+
+ def _get_targets_permutation_indices(self, indices):
+ # permute labels following indices
+ batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+ target_indices = torch.cat([tgt for (_, tgt) in indices])
+ return batch_indices, target_indices
+
+ def forward(
+ self,
+ masks_queries_logits: Tensor,
+ class_queries_logits: Tensor,
+ mask_labels: List[Tensor],
+ class_labels: List[Tensor],
+ auxiliary_predictions: Optional[Dict[str, Tensor]] = None,
+ ) -> Dict[str, Tensor]:
+ """
+ This performs the loss computation.
+
+ Args:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor of shape `batch_size, num_queries, height, width`
+ class_queries_logits (`torch.Tensor`):
+ A tensor of shape `batch_size, num_queries, num_labels`
+ mask_labels (`torch.Tensor`):
+ List of mask labels of shape `(labels, height, width)`.
+ class_labels (`List[torch.Tensor]`):
+ List of class labels of shape `(labels)`.
+ auxiliary_predictions (`Dict[str, torch.Tensor]`, *optional*):
+ if `use_auxiliary_loss` was set to `true` in [`MaskFormerConfig`], then it contains the logits from the
+ inner layers of the Detr's Decoder.
+
+ Returns:
+ `Dict[str, Tensor]`: A dict of `torch.Tensor` containing two keys:
+ - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels.
+ - **loss_mask** -- The loss computed using sigmoid focal loss on the predicted and ground truth masks.
+ - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth
+ masks.
+ if `use_auxiliary_loss` was set to `true` in [`MaskFormerConfig`], the dictionary contains addional losses
+ for each auxiliary predictions.
+ """
+
+ # retrieve the matching between the outputs of the last layer and the labels
+ indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
+ # compute the average number of target masks for normalization purposes
+ num_masks: Number = self.get_num_masks(class_labels, device=class_labels[0].device)
+ # get all the losses
+ losses: Dict[str, Tensor] = {
+ **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks),
+ **self.loss_labels(class_queries_logits, class_labels, indices),
+ }
+ # in case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+ if auxiliary_predictions is not None:
+ for idx, aux_outputs in enumerate(auxiliary_predictions):
+ masks_queries_logits = aux_outputs["masks_queries_logits"]
+ class_queries_logits = aux_outputs["class_queries_logits"]
+ loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
+ loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()}
+ losses.update(loss_dict)
+
+ return losses
+
+ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor:
+ """
+ Computes the average number of target masks across the batch, for normalization purposes.
+ """
+ num_masks = sum([len(classes) for classes in class_labels])
+ num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device)
+ world_size = 1
+ if is_accelerate_available():
+ if PartialState._shared_state != {}:
+ num_masks = reduce(num_masks)
+ world_size = PartialState().num_processes
+
+ num_masks = torch.clamp(num_masks / world_size, min=1)
+ return num_masks
+
+
+class MaskFormerFPNConvLayer(nn.Module):
+ def __init__(self, in_features: int, out_features: int, kernel_size: int = 3, padding: int = 1):
+ """
+ A basic module that executes conv - norm - in sequence used in MaskFormer.
+
+ Args:
+ in_features (`int`):
+ The number of input features (channels).
+ out_features (`int`):
+ The number of outputs features (channels).
+ """
+ super().__init__()
+ self.layers = [
+ nn.Conv2d(in_features, out_features, kernel_size=kernel_size, padding=padding, bias=False),
+ nn.GroupNorm(32, out_features),
+ nn.ReLU(inplace=True),
+ ]
+ for i, layer in enumerate(self.layers):
+ # Provide backwards compatibility from when the class inherited from nn.Sequential
+ # In nn.Sequential subclasses, the name given to the layer is its index in the sequence.
+ # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g.
+ # self.my_layer_name = Layer()
+ # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register
+ # explicitly
+ self.add_module(str(i), layer)
+
+ def forward(self, input: Tensor) -> Tensor:
+ hidden_state = input
+ for layer in self.layers:
+ hidden_state = layer(hidden_state)
+ return hidden_state
+
+
+class MaskFormerFPNLayer(nn.Module):
+ def __init__(self, in_features: int, lateral_features: int):
+ """
+ A Feature Pyramid Network Layer (FPN) layer. It creates a feature map by aggregating features from the previous
+ and backbone layer. Due to the spatial mismatch, the tensor coming from the previous layer is upsampled.
+
+ Args:
+ in_features (`int`):
+ The number of input features (channels).
+ lateral_features (`int`):
+ The number of lateral features (channels).
+ """
+ super().__init__()
+ self.proj = nn.Sequential(
+ nn.Conv2d(lateral_features, in_features, kernel_size=1, padding=0, bias=False),
+ nn.GroupNorm(32, in_features),
+ )
+
+ self.block = MaskFormerFPNConvLayer(in_features, in_features)
+
+ def forward(self, down: Tensor, left: Tensor) -> Tensor:
+ left = self.proj(left)
+ down = nn.functional.interpolate(down, size=left.shape[-2:], mode="nearest")
+ down += left
+ down = self.block(down)
+ return down
+
+
+class MaskFormerFPNModel(nn.Module):
+ def __init__(self, in_features: int, lateral_widths: List[int], feature_size: int = 256):
+ """
+ Feature Pyramid Network, given an input tensor and a set of feature map of different feature/spatial size, it
+ creates a list of feature maps with the same feature size.
+
+ Args:
+ in_features (`int`):
+ The number of input features (channels).
+ lateral_widths (`List[int]`):
+ A list with the features (channels) size of each lateral connection.
+ feature_size (int, *optional*, defaults to 256):
+ The features (channels) of the resulting feature maps.
+ """
+ super().__init__()
+ self.stem = MaskFormerFPNConvLayer(in_features, feature_size)
+ self.layers = nn.Sequential(
+ *[MaskFormerFPNLayer(feature_size, lateral_width) for lateral_width in lateral_widths[::-1]]
+ )
+
+ def forward(self, features: List[Tensor]) -> List[Tensor]:
+ fpn_features = []
+ last_feature = features[-1]
+ other_features = features[:-1]
+ output = self.stem(last_feature)
+ for layer, left in zip(self.layers, other_features[::-1]):
+ output = layer(output, left)
+ fpn_features.append(output)
+ return fpn_features
+
+
+class MaskFormerPixelDecoder(nn.Module):
+ def __init__(self, *args, feature_size: int = 256, mask_feature_size: int = 256, **kwargs):
+ r"""
+ Pixel Decoder Module proposed in [Per-Pixel Classification is Not All You Need for Semantic
+ Segmentation](https://arxiv.org/abs/2107.06278). It first runs the backbone's features into a Feature Pyramid
+ Network creating a list of feature maps. Then, it projects the last one to the correct `mask_size`.
+
+ Args:
+ feature_size (`int`, *optional*, defaults to 256):
+ The feature size (channel dimension) of the FPN feature maps.
+ mask_feature_size (`int`, *optional*, defaults to 256):
+ The features (channels) of the target masks size \\(C_{\epsilon}\\) in the paper.
+ """
+ super().__init__()
+
+ self.fpn = MaskFormerFPNModel(*args, feature_size=feature_size, **kwargs)
+ self.mask_projection = nn.Conv2d(feature_size, mask_feature_size, kernel_size=3, padding=1)
+
+ def forward(
+ self, features: List[Tensor], output_hidden_states: bool = False, return_dict: bool = True
+ ) -> MaskFormerPixelDecoderOutput:
+ fpn_features = self.fpn(features)
+ # we use the last feature map
+ last_feature_projected = self.mask_projection(fpn_features[-1])
+
+ if not return_dict:
+ return (last_feature_projected, tuple(fpn_features)) if output_hidden_states else (last_feature_projected,)
+
+ return MaskFormerPixelDecoderOutput(
+ last_hidden_state=last_feature_projected, hidden_states=tuple(fpn_features) if output_hidden_states else ()
+ )
+
+
+# copied and adapted from original implementation, also practically equal to DetrSinePositionEmbedding
+class MaskFormerSinePositionEmbedding(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, num_pos_feats: int = 64, temperature: int = 10000, normalize: bool = False, scale: Optional[float] = None
+ ):
+ super().__init__()
+ if scale is not None and normalize is False:
+ raise ValueError("normalize should be True if scale is passed")
+ self.num_pos_feats = num_pos_feats
+ self.temperature = temperature
+ self.normalize = normalize
+ self.scale = 2 * math.pi if scale is None else scale
+
+ def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+ if mask is None:
+ mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
+ not_mask = (~mask).to(x.dtype)
+ y_embed = not_mask.cumsum(1)
+ x_embed = not_mask.cumsum(2)
+ if self.normalize:
+ eps = 1e-6
+ y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+ x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+ dim_t = torch.arange(self.num_pos_feats, dtype=torch.int64, device=x.device).type_as(x)
+ dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_pos_feats)
+
+ 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 PredictionBlock(nn.Module):
+ def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None:
+ super().__init__()
+ self.layers = [nn.Linear(in_dim, out_dim), activation]
+ # Maintain submodule indexing as if part of a Sequential block
+ for i, layer in enumerate(self.layers):
+ self.add_module(str(i), layer)
+
+ def forward(self, input: Tensor) -> Tensor:
+ hidden_state = input
+ for layer in self.layers:
+ hidden_state = layer(hidden_state)
+ return hidden_state
+
+
+class MaskformerMLPPredictionHead(nn.Module):
+ def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3):
+ """
+ A classic Multi Layer Perceptron (MLP).
+
+ Args:
+ input_dim (`int`):
+ The input dimensions.
+ hidden_dim (`int`):
+ The hidden dimensions.
+ output_dim (`int`):
+ The output dimensions.
+ num_layers (int, *optional*, defaults to 3):
+ The number of layers.
+ """
+ super().__init__()
+ in_dims = [input_dim] + [hidden_dim] * (num_layers - 1)
+ out_dims = [hidden_dim] * (num_layers - 1) + [output_dim]
+
+ self.layers = []
+ for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)):
+ activation = nn.ReLU() if i < num_layers - 1 else nn.Identity()
+ layer = PredictionBlock(in_dim, out_dim, activation=activation)
+ self.layers.append(layer)
+ # Provide backwards compatibility from when the class inherited from nn.Sequential
+ # In nn.Sequential subclasses, the name given to the layer is its index in the sequence.
+ # In nn.Module subclasses they derived from the instance attribute they are assigned to e.g.
+ # self.my_layer_name = Layer()
+ # We can't give instance attributes integer names i.e. self.0 is not permitted and so need to register
+ # explicitly
+ self.add_module(str(i), layer)
+
+ def forward(self, input: Tensor) -> Tensor:
+ hidden_state = input
+ for layer in self.layers:
+ hidden_state = layer(hidden_state)
+ return hidden_state
+
+
+class MaskFormerPixelLevelModule(nn.Module):
+ def __init__(self, config: MaskFormerConfig):
+ """
+ Pixel Level Module proposed in [Per-Pixel Classification is Not All You Need for Semantic
+ Segmentation](https://arxiv.org/abs/2107.06278). It runs the input image through a backbone and a pixel
+ decoder, generating an image feature map and pixel embeddings.
+
+ Args:
+ config ([`MaskFormerConfig`]):
+ The configuration used to instantiate this model.
+ """
+ super().__init__()
+ if getattr(config, "backbone_config") is not None and config.backbone_config.model_type == "swin":
+ # for backwards compatibility
+ backbone_config = config.backbone_config
+ backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict())
+ backbone_config.out_features = ["stage1", "stage2", "stage3", "stage4"]
+ config.backbone_config = backbone_config
+ self.encoder = load_backbone(config)
+
+ feature_channels = self.encoder.channels
+ self.decoder = MaskFormerPixelDecoder(
+ in_features=feature_channels[-1],
+ feature_size=config.fpn_feature_size,
+ mask_feature_size=config.mask_feature_size,
+ lateral_widths=feature_channels[:-1],
+ )
+
+ def forward(
+ self, pixel_values: Tensor, output_hidden_states: bool = False, return_dict: bool = True
+ ) -> MaskFormerPixelLevelModuleOutput:
+ features = self.encoder(pixel_values).feature_maps
+ decoder_output = self.decoder(features, output_hidden_states, return_dict=return_dict)
+
+ if not return_dict:
+ last_hidden_state = decoder_output[0]
+ outputs = (features[-1], last_hidden_state)
+ if output_hidden_states:
+ hidden_states = decoder_output[1]
+ outputs = outputs + (tuple(features),) + (hidden_states,)
+ return outputs
+
+ return MaskFormerPixelLevelModuleOutput(
+ # the last feature is actually the output from the last layer
+ encoder_last_hidden_state=features[-1],
+ decoder_last_hidden_state=decoder_output.last_hidden_state,
+ encoder_hidden_states=tuple(features) if output_hidden_states else (),
+ decoder_hidden_states=decoder_output.hidden_states if output_hidden_states else (),
+ )
+
+
+class MaskFormerTransformerModule(nn.Module):
+ """
+ The MaskFormer's transformer module.
+ """
+
+ def __init__(self, in_features: int, config: MaskFormerConfig):
+ super().__init__()
+ hidden_size = config.decoder_config.hidden_size
+ should_project = in_features != hidden_size
+ self.position_embedder = MaskFormerSinePositionEmbedding(num_pos_feats=hidden_size // 2, normalize=True)
+ self.queries_embedder = nn.Embedding(config.decoder_config.num_queries, hidden_size)
+ self.input_projection = nn.Conv2d(in_features, hidden_size, kernel_size=1) if should_project else None
+ self.decoder = DetrDecoder(config=config.decoder_config)
+
+ def forward(
+ self,
+ image_features: Tensor,
+ output_hidden_states: bool = False,
+ output_attentions: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> DetrDecoderOutput:
+ if self.input_projection is not None:
+ image_features = self.input_projection(image_features)
+ object_queries = self.position_embedder(image_features)
+ # repeat the queries "q c -> b q c"
+ batch_size = image_features.shape[0]
+ queries_embeddings = self.queries_embedder.weight.unsqueeze(0).repeat(batch_size, 1, 1)
+ inputs_embeds = torch.zeros_like(queries_embeddings, requires_grad=self.training)
+
+ # torch.export.export does no support requires_grad
+ if self.training:
+ inputs_embeds.requires_grad_(True)
+
+ batch_size, num_channels, height, width = image_features.shape
+ # rearrange both image_features and object_queries "b c h w -> b (h w) c"
+ image_features = image_features.view(batch_size, num_channels, height * width).permute(0, 2, 1)
+ object_queries = object_queries.view(batch_size, num_channels, height * width).permute(0, 2, 1)
+
+ decoder_output: DetrDecoderOutput = self.decoder(
+ inputs_embeds=inputs_embeds,
+ attention_mask=None,
+ encoder_hidden_states=image_features,
+ encoder_attention_mask=None,
+ object_queries=object_queries,
+ query_position_embeddings=queries_embeddings,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ return decoder_output
+
+
+MASKFORMER_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 ([`MaskFormerConfig`]): 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.
+"""
+
+MASKFORMER_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
+ [`MaskFormerImageProcessor.__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)
+ 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_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of Detr's decoder attention layers.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~MaskFormerModelOutput`] instead of a plain tuple.
+"""
+
+
+class MaskFormerPreTrainedModel(PreTrainedModel):
+ config_class = MaskFormerConfig
+ base_model_prefix = "model"
+ main_input_name = "pixel_values"
+
+ def _init_weights(self, module: nn.Module):
+ xavier_std = self.config.init_xavier_std
+ std = self.config.init_std
+ if isinstance(module, MaskFormerTransformerModule):
+ if module.input_projection is not None:
+ nn.init.xavier_uniform_(module.input_projection.weight, gain=xavier_std)
+ nn.init.constant_(module.input_projection.bias, 0)
+ # FPN
+ elif isinstance(module, MaskFormerFPNModel):
+ nn.init.xavier_uniform_(module.stem.get_submodule("0").weight, gain=xavier_std)
+
+ elif isinstance(module, MaskFormerFPNLayer):
+ nn.init.xavier_uniform_(module.proj[0].weight, gain=xavier_std)
+
+ elif isinstance(module, MaskFormerFPNConvLayer):
+ nn.init.xavier_uniform_(module.get_submodule("0").weight, gain=xavier_std)
+ # The MLP head
+ elif isinstance(module, MaskformerMLPPredictionHead):
+ # I was not able to find the correct initializer in the original implementation
+ # we'll use xavier
+ for submodule in module.modules():
+ if isinstance(submodule, nn.Linear):
+ nn.init.xavier_uniform_(submodule.weight, gain=xavier_std)
+ nn.init.constant_(submodule.bias, 0)
+ elif isinstance(module, nn.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+ # copied from DETR
+ 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_()
+
+
+@add_start_docstrings(
+ "The bare MaskFormer Model outputting raw hidden-states without any specific head on top.",
+ MASKFORMER_START_DOCSTRING,
+)
+class MaskFormerModel(MaskFormerPreTrainedModel):
+ def __init__(self, config: MaskFormerConfig):
+ super().__init__(config)
+ self.pixel_level_module = MaskFormerPixelLevelModule(config)
+ self.transformer_module = MaskFormerTransformerModule(
+ in_features=self.pixel_level_module.encoder.channels[-1], config=config
+ )
+
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(MASKFORMER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=MaskFormerModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ pixel_values: Tensor,
+ pixel_mask: Optional[Tensor] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> MaskFormerModelOutput:
+ r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, MaskFormerModel
+ >>> from PIL import Image
+ >>> import requests
+
+ >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation
+ >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade")
+ >>> model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-base-ade")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = image_processor(image, return_tensors="pt")
+
+ >>> # forward pass
+ >>> outputs = model(**inputs)
+
+ >>> # the decoder of MaskFormer outputs hidden states of shape (batch_size, num_queries, hidden_size)
+ >>> transformer_decoder_last_hidden_state = outputs.transformer_decoder_last_hidden_state
+ >>> list(transformer_decoder_last_hidden_state.shape)
+ [1, 100, 256]
+ ```"""
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_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
+
+ batch_size, _, height, width = pixel_values.shape
+
+ if pixel_mask is None:
+ pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device)
+
+ pixel_level_module_output = self.pixel_level_module(
+ pixel_values, output_hidden_states, return_dict=return_dict
+ )
+ image_features = pixel_level_module_output[0]
+ pixel_embeddings = pixel_level_module_output[1]
+
+ transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions)
+ queries = transformer_module_output.last_hidden_state
+
+ encoder_hidden_states = None
+ pixel_decoder_hidden_states = None
+ transformer_decoder_hidden_states = None
+ hidden_states = None
+
+ if output_hidden_states:
+ encoder_hidden_states = pixel_level_module_output[2]
+ pixel_decoder_hidden_states = pixel_level_module_output[3]
+ transformer_decoder_hidden_states = transformer_module_output[1]
+ hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states
+
+ output = MaskFormerModelOutput(
+ encoder_last_hidden_state=image_features,
+ pixel_decoder_last_hidden_state=pixel_embeddings,
+ transformer_decoder_last_hidden_state=queries,
+ encoder_hidden_states=encoder_hidden_states,
+ pixel_decoder_hidden_states=pixel_decoder_hidden_states,
+ transformer_decoder_hidden_states=transformer_decoder_hidden_states,
+ hidden_states=hidden_states,
+ attentions=transformer_module_output.attentions,
+ )
+
+ if not return_dict:
+ output = tuple(v for v in output.values())
+
+ return output
+
+
+class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel):
+ def __init__(self, config: MaskFormerConfig):
+ super().__init__(config)
+ self.model = MaskFormerModel(config)
+ hidden_size = config.decoder_config.hidden_size
+ # + 1 because we add the "null" class
+ self.class_predictor = nn.Linear(hidden_size, config.num_labels + 1)
+ self.mask_embedder = MaskformerMLPPredictionHead(hidden_size, hidden_size, config.mask_feature_size)
+
+ self.matcher = MaskFormerHungarianMatcher(
+ cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight
+ )
+
+ self.weight_dict: Dict[str, float] = {
+ "loss_cross_entropy": config.cross_entropy_weight,
+ "loss_mask": config.mask_weight,
+ "loss_dice": config.dice_weight,
+ }
+
+ self.criterion = MaskFormerLoss(
+ config.num_labels,
+ matcher=self.matcher,
+ weight_dict=self.weight_dict,
+ eos_coef=config.no_object_weight,
+ )
+
+ self.post_init()
+
+ def get_loss_dict(
+ self,
+ masks_queries_logits: Tensor,
+ class_queries_logits: Tensor,
+ mask_labels: Tensor,
+ class_labels: Tensor,
+ auxiliary_logits: Dict[str, Tensor],
+ ) -> Dict[str, Tensor]:
+ loss_dict: Dict[str, Tensor] = self.criterion(
+ masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits
+ )
+ # weight each loss by `self.weight_dict[]` including auxiliary losses
+ for key, weight in self.weight_dict.items():
+ for loss_key, loss in loss_dict.items():
+ if key in loss_key:
+ loss *= weight
+
+ return loss_dict
+
+ def get_loss(self, loss_dict: Dict[str, Tensor]) -> Tensor:
+ return sum(loss_dict.values())
+
+ def get_logits(self, outputs: MaskFormerModelOutput) -> Tuple[Tensor, Tensor, Dict[str, Tensor]]:
+ pixel_embeddings = outputs.pixel_decoder_last_hidden_state
+ # get the auxiliary predictions (one for each decoder's layer)
+ auxiliary_logits: List[str, Tensor] = []
+
+ is_tracing = torch.jit.is_tracing() or isinstance(outputs, torch.fx.Proxy) or is_torchdynamo_compiling()
+ # This code is a little bit cumbersome, an improvement can be to return a list of predictions. If we have auxiliary loss then we are going to return more than one element in the list
+ if self.config.use_auxiliary_loss:
+ stacked_transformer_decoder_outputs = torch.stack(outputs.transformer_decoder_hidden_states)
+ classes = self.class_predictor(stacked_transformer_decoder_outputs)
+ class_queries_logits = classes[-1]
+ # get the masks
+ mask_embeddings = self.mask_embedder(stacked_transformer_decoder_outputs)
+
+ if is_tracing and not is_torch_greater_or_equal_than_2_1:
+ # Equivalent to einsum('lbqc, bchw -> lbqhw') but jit friendly
+ num_embeddings, batch_size, num_queries, num_channels = mask_embeddings.shape
+ _, _, height, width = pixel_embeddings.shape
+ binaries_masks = torch.zeros(
+ (num_embeddings, batch_size, num_queries, height, width), device=mask_embeddings.device
+ )
+ for c in range(num_channels):
+ binaries_masks += mask_embeddings[..., c][..., None, None] * pixel_embeddings[None, :, None, c]
+ else:
+ binaries_masks = torch.einsum("lbqc, bchw -> lbqhw", mask_embeddings, pixel_embeddings)
+
+ masks_queries_logits = binaries_masks[-1]
+ # go til [:-1] because the last one is always used
+ for aux_binary_masks, aux_classes in zip(binaries_masks[:-1], classes[:-1]):
+ auxiliary_logits.append(
+ {"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes}
+ )
+
+ else:
+ transformer_decoder_hidden_states = outputs.transformer_decoder_last_hidden_state
+ classes = self.class_predictor(transformer_decoder_hidden_states)
+ class_queries_logits = classes
+ # get the masks
+ mask_embeddings = self.mask_embedder(transformer_decoder_hidden_states)
+ # sum up over the channels
+
+ if is_tracing and not is_torch_greater_or_equal_than_2_1:
+ # Equivalent to einsum('bqc, bchw -> bqhw') but jit friendly
+ batch_size, num_queries, num_channels = mask_embeddings.shape
+ _, _, height, width = pixel_embeddings.shape
+ masks_queries_logits = torch.zeros(
+ (batch_size, num_queries, height, width), device=mask_embeddings.device
+ )
+ for c in range(num_channels):
+ masks_queries_logits += mask_embeddings[..., c][..., None, None] * pixel_embeddings[:, None, c]
+ else:
+ masks_queries_logits = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings)
+
+ return class_queries_logits, masks_queries_logits, auxiliary_logits
+
+ @add_start_docstrings_to_model_forward(MASKFORMER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=MaskFormerForInstanceSegmentationOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ pixel_values: Tensor,
+ mask_labels: Optional[List[Tensor]] = None,
+ class_labels: Optional[List[Tensor]] = None,
+ pixel_mask: Optional[Tensor] = None,
+ output_auxiliary_logits: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> MaskFormerForInstanceSegmentationOutput:
+ r"""
+ mask_labels (`List[torch.Tensor]`, *optional*):
+ List of mask labels of shape `(num_labels, height, width)` to be fed to a model
+ class_labels (`List[torch.LongTensor]`, *optional*):
+ list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the
+ labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`.
+
+ Returns:
+
+ Examples:
+
+ Semantic segmentation example:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation
+ >>> from PIL import Image
+ >>> import requests
+
+ >>> # load MaskFormer fine-tuned on ADE20k semantic segmentation
+ >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-ade")
+ >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-ade")
+
+ >>> url = (
+ ... "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
+ ... )
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)`
+ >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
+ >>> class_queries_logits = outputs.class_queries_logits
+ >>> masks_queries_logits = outputs.masks_queries_logits
+
+ >>> # you can pass them to image_processor for postprocessing
+ >>> predicted_semantic_map = image_processor.post_process_semantic_segmentation(
+ ... outputs, target_sizes=[(image.height, image.width)]
+ ... )[0]
+
+ >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs)
+ >>> list(predicted_semantic_map.shape)
+ [512, 683]
+ ```
+
+ Panoptic segmentation example:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, MaskFormerForInstanceSegmentation
+ >>> from PIL import Image
+ >>> import requests
+
+ >>> # load MaskFormer fine-tuned on COCO panoptic segmentation
+ >>> image_processor = AutoImageProcessor.from_pretrained("facebook/maskformer-swin-base-coco")
+ >>> model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-base-coco")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> # model predicts class_queries_logits of shape `(batch_size, num_queries)`
+ >>> # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
+ >>> class_queries_logits = outputs.class_queries_logits
+ >>> masks_queries_logits = outputs.masks_queries_logits
+
+ >>> # you can pass them to image_processor for postprocessing
+ >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(image.height, image.width)])[0]
+
+ >>> # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs)
+ >>> predicted_panoptic_map = result["segmentation"]
+ >>> list(predicted_panoptic_map.shape)
+ [480, 640]
+ ```
+ """
+
+ 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
+
+ raw_outputs = self.model(
+ pixel_values,
+ pixel_mask,
+ output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss,
+ return_dict=return_dict,
+ output_attentions=output_attentions,
+ )
+ # We need to have raw_outputs optionally be returned as a dict to use torch.compile. For backwards
+ # compatibility we convert to a dataclass for the rest of the model logic
+ outputs = MaskFormerModelOutput(
+ encoder_last_hidden_state=raw_outputs[0],
+ pixel_decoder_last_hidden_state=raw_outputs[1],
+ transformer_decoder_last_hidden_state=raw_outputs[2],
+ encoder_hidden_states=raw_outputs[3] if output_hidden_states else None,
+ pixel_decoder_hidden_states=raw_outputs[4] if output_hidden_states else None,
+ transformer_decoder_hidden_states=raw_outputs[5] if output_hidden_states else None,
+ hidden_states=raw_outputs[6] if output_hidden_states else None,
+ attentions=raw_outputs[-1] if output_attentions else None,
+ )
+
+ loss, loss_dict, auxiliary_logits = None, None, None
+
+ class_queries_logits, masks_queries_logits, auxiliary_logits = self.get_logits(outputs)
+
+ if mask_labels is not None and class_labels is not None:
+ loss_dict: Dict[str, Tensor] = self.get_loss_dict(
+ masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits
+ )
+ loss = self.get_loss(loss_dict)
+
+ output_auxiliary_logits = (
+ self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits
+ )
+ if not output_auxiliary_logits:
+ auxiliary_logits = None
+
+ if not return_dict:
+ output = tuple(
+ v
+ for v in (loss, class_queries_logits, masks_queries_logits, auxiliary_logits, *outputs.values())
+ if v is not None
+ )
+ return output
+
+ return MaskFormerForInstanceSegmentationOutput(
+ loss=loss,
+ **outputs,
+ class_queries_logits=class_queries_logits,
+ masks_queries_logits=masks_queries_logits,
+ auxiliary_logits=auxiliary_logits,
+ )
+
+
+__all__ = ["MaskFormerForInstanceSegmentation", "MaskFormerModel", "MaskFormerPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer_swin.py b/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd3fc11ca1650efa248ed2a07fbae06765058e05
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/maskformer/modeling_maskformer_swin.py
@@ -0,0 +1,965 @@
+# 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.
+
+"""MaskFormer Swin Transformer. The reason Swin Transformer is implemented here is because MaskFormer uses the hidden
+states before downsampling, which is different from the default Swin Transformer."""
+
+import collections.abc
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from ...activations import ACT2FN
+from ...file_utils import ModelOutput
+from ...modeling_outputs import BackboneOutput
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
+from ...utils import torch_int
+from ...utils.backbone_utils import BackboneMixin
+from .configuration_maskformer_swin import MaskFormerSwinConfig
+
+
+@dataclass
+class MaskFormerSwinModelOutputWithPooling(ModelOutput):
+ """
+ Class for MaskFormerSwinModel's outputs that also contains the spatial dimensions of the 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)`):
+ Last layer hidden-state after a mean pooling operation.
+ 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.
+ hidden_states_spatial_dimensions (`tuple(tuple(int, int))`, *optional*):
+ A tuple containing the spatial dimension of each `hidden_state` needed to reshape the `hidden_states` to
+ `batch, channels, height, width`. Due to padding, their spatial size cannot be inferred before the
+ `forward` method.
+ 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: Optional[torch.FloatTensor] = None
+ pooler_output: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ hidden_states_spatial_dimensions: Tuple[Tuple[int, int]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class MaskFormerSwinBaseModelOutput(ModelOutput):
+ """
+ Class for SwinEncoder's outputs.
+
+ 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.
+ hidden_states_spatial_dimensions (`tuple(tuple(int, int))`, *optional*):
+ A tuple containing the spatial dimension of each `hidden_state` needed to reshape the `hidden_states` to
+ `batch, channels, height, width`. Due to padding, their spatial size cannot inferred before the `forward`
+ method.
+ 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: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ hidden_states_spatial_dimensions: Tuple[Tuple[int, int]] = None
+ attentions: 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.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
+
+
+class MaskFormerSwinEmbeddings(nn.Module):
+ """
+ Construct the patch and position embeddings.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+
+ self.patch_embeddings = MaskFormerSwinPatchEmbeddings(config)
+ num_patches = self.patch_embeddings.num_patches
+ self.patch_grid = self.patch_embeddings.grid_size
+
+ 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)
+ self.patch_size = config.patch_size
+
+ # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding
+ 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. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embeddings.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, :1]
+ patch_pos_embed = self.position_embeddings[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values, interpolate_pos_encoding):
+ _, num_channels, height, width = pixel_values.shape
+ embeddings, output_dimensions = self.patch_embeddings(pixel_values)
+ embeddings = self.norm(embeddings)
+
+ if self.position_embeddings is not None:
+ 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, output_dimensions
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinPatchEmbeddings with Swin->MaskFormerSwin
+class MaskFormerSwinPatchEmbeddings(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
+ # 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 MaskFormerSwinPatchMerging(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.swin.modeling_swin.SwinDropPath with Swin->MaskFormerSwin
+class MaskFormerSwinDropPath(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->MaskFormerSwin
+class MaskFormerSwinSelfAttention(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 MaskFormerSwinModel 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 with Swin->MaskFormerSwin
+class MaskFormerSwinSelfOutput(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->MaskFormerSwin
+class MaskFormerSwinAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size):
+ super().__init__()
+ self.self = MaskFormerSwinSelfAttention(config, dim, num_heads, window_size)
+ self.output = MaskFormerSwinSelfOutput(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 with Swin->MaskFormerSwin
+class MaskFormerSwinIntermediate(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 with Swin->MaskFormerSwin
+class MaskFormerSwinOutput(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 MaskFormerSwinLayer(nn.Module):
+ def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0):
+ super().__init__()
+ 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 = MaskFormerSwinAttention(config, dim, num_heads, self.window_size)
+ self.drop_path = MaskFormerSwinDropPath(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 = MaskFormerSwinIntermediate(config, dim)
+ self.output = MaskFormerSwinOutput(config, dim)
+
+ def get_attn_mask(self, input_resolution):
+ if self.shift_size > 0:
+ # calculate attention mask for SW-MSA
+ height, width = input_resolution
+ img_mask = torch.zeros((1, height, width, 1))
+ 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_left = pad_top = 0
+ pad_rigth = (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, pad_left, pad_rigth, pad_top, pad_bottom)
+ hidden_states = nn.functional.pad(hidden_states, pad_values)
+ return hidden_states, pad_values
+
+ def forward(self, hidden_states, input_dimensions, head_mask=None, output_attentions=False):
+ height, width = input_dimensions
+ batch_size, dim, 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))
+ if attn_mask is not None:
+ attn_mask = attn_mask.to(hidden_states_windows.device)
+
+ self_attention_outputs = self.attention(
+ hidden_states_windows, attn_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
+
+ 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
+ ) # B height' width' C
+
+ # 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)
+
+ outputs = (layer_output,) + outputs
+
+ return outputs
+
+
+class MaskFormerSwinStage(nn.Module):
+ # Copied from transformers.models.swin.modeling_swin.SwinStage.__init__ with Swin->MaskFormerSwin
+ 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(
+ [
+ MaskFormerSwinLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ 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, input_dimensions, head_mask=None, output_attentions=False, output_hidden_states=False
+ ):
+ all_hidden_states = () if output_hidden_states else None
+
+ height, width = input_dimensions
+ for i, block_module in enumerate(self.blocks):
+ 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
+
+ block_hidden_states = block_module(hidden_states, input_dimensions, layer_head_mask, output_attentions)
+
+ hidden_states = block_hidden_states[0]
+
+ if output_hidden_states:
+ all_hidden_states += (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, input_dimensions)
+ else:
+ output_dimensions = (height, width, height, width)
+
+ return hidden_states, output_dimensions, all_hidden_states
+
+
+class MaskFormerSwinEncoder(nn.Module):
+ # Copied from transformers.models.swin.modeling_swin.SwinEncoder.__init__ with Swin->MaskFormerSwin
+ 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(
+ [
+ MaskFormerSwinStage(
+ 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=MaskFormerSwinPatchMerging 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,
+ input_dimensions,
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ ):
+ all_hidden_states = () if output_hidden_states else None
+ all_input_dimensions = ()
+ all_self_attentions = () if output_attentions else None
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ 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_hidden_states, output_dimensions, layer_all_hidden_states = self._gradient_checkpointing_func(
+ layer_module.__call__,
+ hidden_states,
+ layer_head_mask,
+ output_attentions,
+ )
+ else:
+ layer_hidden_states, output_dimensions, layer_all_hidden_states = layer_module(
+ hidden_states,
+ input_dimensions,
+ layer_head_mask,
+ output_attentions,
+ output_hidden_states,
+ )
+
+ input_dimensions = (output_dimensions[-2], output_dimensions[-1])
+ all_input_dimensions += (input_dimensions,)
+ if output_hidden_states:
+ all_hidden_states += (layer_all_hidden_states,)
+
+ hidden_states = layer_hidden_states
+
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_all_hidden_states[1],)
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+ return MaskFormerSwinBaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ hidden_states_spatial_dimensions=all_input_dimensions,
+ attentions=all_self_attentions,
+ )
+
+
+class MaskFormerSwinPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MaskFormerSwinConfig
+ base_model_prefix = "model"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MaskFormerSwinStage"]
+
+ 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)
+ elif isinstance(module, MaskFormerSwinEmbeddings):
+ if module.position_embeddings is not None:
+ module.position_embeddings.data.zero_()
+ elif isinstance(module, MaskFormerSwinSelfAttention):
+ module.relative_position_bias_table.data.zero_()
+
+
+class MaskFormerSwinModel(MaskFormerSwinPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=True):
+ 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 = MaskFormerSwinEmbeddings(config)
+ self.encoder = MaskFormerSwinEncoder(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
+
+ 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)
+
+ def forward(
+ self,
+ pixel_values=None,
+ head_mask=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ interpolate_pos_encoding=False,
+ return_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.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, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ 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.last_hidden_state if return_dict else 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:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ hidden_states_spatial_dimensions = (input_dimensions,) + encoder_outputs.hidden_states_spatial_dimensions
+
+ return MaskFormerSwinModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ hidden_states_spatial_dimensions=hidden_states_spatial_dimensions,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+class MaskFormerSwinBackbone(MaskFormerSwinPreTrainedModel, BackboneMixin):
+ """
+ MaskFormerSwin backbone, designed especially for the MaskFormer framework.
+
+ This classes reshapes `hidden_states` from (`batch_size, sequence_length, hidden_size)` to (`batch_size,
+ num_channels, height, width)`). It also adds additional layernorms after each stage.
+
+ Args:
+ config (`MaskFormerSwinConfig`):
+ The configuration used by [`MaskFormerSwinModel`].
+ """
+
+ def __init__(self, config: MaskFormerSwinConfig):
+ super().__init__(config)
+ super()._init_backbone(config)
+
+ self.model = MaskFormerSwinModel(config)
+ if "stem" in self.out_features:
+ raise ValueError("This backbone does not support 'stem' in the `out_features`.")
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+ self.hidden_states_norms = nn.ModuleList(
+ [nn.LayerNorm(num_channels) for num_channels in self.num_features[1:]]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def forward(
+ self,
+ pixel_values: Tensor,
+ output_hidden_states: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> BackboneOutput:
+ 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.model(
+ pixel_values, output_hidden_states=True, output_attentions=output_attentions, return_dict=True
+ )
+
+ # we skip the stem
+ hidden_states = outputs.hidden_states[1:]
+
+ # we need to reshape the hidden states to their original spatial dimensions
+ # spatial dimensions contains all the heights and widths of each stage, including after the embeddings
+ spatial_dimensions: Tuple[Tuple[int, int]] = outputs.hidden_states_spatial_dimensions
+ feature_maps = ()
+ for i, (hidden_state, stage, (height, width)) in enumerate(
+ zip(hidden_states, self.stage_names[1:], spatial_dimensions)
+ ):
+ norm = self.hidden_states_norms[i]
+ # the last element corespond to the layer's last block output but before patch merging
+ hidden_state_unpolled = hidden_state[-1]
+ hidden_state_norm = norm(hidden_state_unpolled)
+ # the pixel decoder (FPN) expects 3D tensors (features)
+ batch_size, _, hidden_size = hidden_state_norm.shape
+ # reshape "b (h w) d -> b d h w"
+ hidden_state_permuted = (
+ hidden_state_norm.permute(0, 2, 1).view((batch_size, hidden_size, height, width)).contiguous()
+ )
+ if stage in self.out_features:
+ feature_maps += (hidden_state_permuted,)
+
+ if not return_dict:
+ output = (feature_maps,)
+ if output_hidden_states:
+ output += (outputs.hidden_states,)
+ if output_attentions:
+ output += (outputs.attentions,)
+ return output
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.hidden_states if output_hidden_states else None,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["MaskFormerSwinBackbone", "MaskFormerSwinModel", "MaskFormerSwinPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__init__.py b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..05c7e8615ddf0f8bac681b0ed6e124c1f1a9a30b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__init__.py
@@ -0,0 +1,29 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_mobilenet_v1 import *
+ from .feature_extraction_mobilenet_v1 import *
+ from .image_processing_mobilenet_v1 import *
+ from .modeling_mobilenet_v1 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0953975a6d7a6066244aa0c43c9f8bf299049a1f
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/configuration_mobilenet_v1.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/configuration_mobilenet_v1.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ee809494175bd20fcb611af23e99699598614307
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/configuration_mobilenet_v1.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/feature_extraction_mobilenet_v1.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/feature_extraction_mobilenet_v1.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..31fbe99bf59a5f9e2727027c49524d185afde377
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/feature_extraction_mobilenet_v1.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/image_processing_mobilenet_v1.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/image_processing_mobilenet_v1.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3284368f52bc1329f77957a1e33ae7421a84ffbe
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/image_processing_mobilenet_v1.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/modeling_mobilenet_v1.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/modeling_mobilenet_v1.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7b5a6f3a644f4b0bb8df17aefdee73e1afa1875e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/__pycache__/modeling_mobilenet_v1.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/configuration_mobilenet_v1.py b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/configuration_mobilenet_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..59e977b95212756a0e7bf061e251e41dbd44d020
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/configuration_mobilenet_v1.py
@@ -0,0 +1,126 @@
+# 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.
+"""MobileNetV1 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__)
+
+
+class MobileNetV1Config(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MobileNetV1Model`]. It is used to instantiate a
+ MobileNetV1 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 MobileNetV1
+ [google/mobilenet_v1_1.0_224](https://huggingface.co/google/mobilenet_v1_1.0_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.
+ image_size (`int`, *optional*, defaults to 224):
+ The size (resolution) of each image.
+ depth_multiplier (`float`, *optional*, defaults to 1.0):
+ Shrinks or expands the number of channels in each layer. Default is 1.0, which starts the network with 32
+ channels. This is sometimes also called "alpha" or "width multiplier".
+ min_depth (`int`, *optional*, defaults to 8):
+ All layers will have at least this many channels.
+ hidden_act (`str` or `function`, *optional*, defaults to `"relu6"`):
+ The non-linear activation function (function or string) in the Transformer encoder and convolution layers.
+ tf_padding (`bool`, *optional*, defaults to `True`):
+ Whether to use TensorFlow padding rules on the convolution layers.
+ classifier_dropout_prob (`float`, *optional*, defaults to 0.999):
+ The dropout ratio for attached classifiers.
+ 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 0.001):
+ The epsilon used by the layer normalization layers.
+
+ Example:
+
+ ```python
+ >>> from transformers import MobileNetV1Config, MobileNetV1Model
+
+ >>> # Initializing a "mobilenet_v1_1.0_224" style configuration
+ >>> configuration = MobileNetV1Config()
+
+ >>> # Initializing a model from the "mobilenet_v1_1.0_224" style configuration
+ >>> model = MobileNetV1Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "mobilenet_v1"
+
+ def __init__(
+ self,
+ num_channels=3,
+ image_size=224,
+ depth_multiplier=1.0,
+ min_depth=8,
+ hidden_act="relu6",
+ tf_padding=True,
+ classifier_dropout_prob=0.999,
+ initializer_range=0.02,
+ layer_norm_eps=0.001,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ if depth_multiplier <= 0:
+ raise ValueError("depth_multiplier must be greater than zero.")
+
+ self.num_channels = num_channels
+ self.image_size = image_size
+ self.depth_multiplier = depth_multiplier
+ self.min_depth = min_depth
+ self.hidden_act = hidden_act
+ self.tf_padding = tf_padding
+ self.classifier_dropout_prob = classifier_dropout_prob
+ self.initializer_range = initializer_range
+ self.layer_norm_eps = layer_norm_eps
+
+
+class MobileNetV1OnnxConfig(OnnxConfig):
+ torch_onnx_minimum_version = version.parse("1.11")
+
+ @property
+ def inputs(self) -> Mapping[str, Mapping[int, str]]:
+ return OrderedDict([("pixel_values", {0: "batch"})])
+
+ @property
+ def outputs(self) -> Mapping[str, Mapping[int, str]]:
+ if self.task == "image-classification":
+ return OrderedDict([("logits", {0: "batch"})])
+ else:
+ return OrderedDict([("last_hidden_state", {0: "batch"}), ("pooler_output", {0: "batch"})])
+
+ @property
+ def atol_for_validation(self) -> float:
+ return 1e-4
+
+
+__all__ = ["MobileNetV1Config", "MobileNetV1OnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/feature_extraction_mobilenet_v1.py b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/feature_extraction_mobilenet_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0d2d11a207241565666c9382ab786a8af9255c7
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/feature_extraction_mobilenet_v1.py
@@ -0,0 +1,36 @@
+# 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 MobileNetV1."""
+
+import warnings
+
+from ...utils import logging
+from .image_processing_mobilenet_v1 import MobileNetV1ImageProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+class MobileNetV1FeatureExtractor(MobileNetV1ImageProcessor):
+ def __init__(self, *args, **kwargs) -> None:
+ warnings.warn(
+ "The class MobileNetV1FeatureExtractor is deprecated and will be removed in version 5 of Transformers."
+ " Please use MobileNetV1ImageProcessor instead.",
+ FutureWarning,
+ )
+ super().__init__(*args, **kwargs)
+
+
+__all__ = ["MobileNetV1FeatureExtractor"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/image_processing_mobilenet_v1.py b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/image_processing_mobilenet_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c30c3413baa58b34a50308725b02acfb6d8ac01
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/image_processing_mobilenet_v1.py
@@ -0,0 +1,305 @@
+# 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 MobileNetV1."""
+
+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_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_preprocess_arguments,
+)
+from ...utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MobileNetV1ImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a MobileNetV1 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": 256}`):
+ 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 `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. 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 the `do_center_crop` parameter in the
+ `preprocess` method.
+ crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Desired output size when applying center-cropping. Only has an effect if `do_center_crop` is set to `True`.
+ Can be overridden by the `crop_size` 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.
+ 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.BILINEAR,
+ 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, 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": 256}
+ size = get_size_dict(size, default_to_square=False)
+ crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
+ crop_size = get_size_dict(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
+
+ # Copied from transformers.models.clip.image_processing_clip.CLIPImageProcessor.resize
+ 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. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
+ resized to keep the input aspect ratio.
+
+ 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.
+ """
+ default_to_square = True
+ if "shortest_edge" in size:
+ size = size["shortest_edge"]
+ default_to_square = False
+ elif "height" in size and "width" in size:
+ size = (size["height"], size["width"])
+ else:
+ raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
+
+ output_size = get_resize_output_image_size(
+ image,
+ size=size,
+ default_to_square=default_to_square,
+ 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,
+ )
+
+ @filter_out_non_signature_kwargs()
+ 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: 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,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ):
+ """
+ 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 size["shortest_edge"], with
+ the longest edge resized to keep the input aspect ratio.
+ 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.
+ 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 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
+ size = size if size is not None else self.size
+ size = get_size_dict(size, default_to_square=False)
+ 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)
+ 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)
+
+ 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 do_rescale and is_scaled_image(images[0]):
+ 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])
+
+ all_images = []
+ for image in images:
+ 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
+ )
+
+ all_images.append(image)
+ images = [
+ to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+ for image in all_images
+ ]
+
+ data = {"pixel_values": images}
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MobileNetV1ImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilenet_v1/modeling_mobilenet_v1.py b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/modeling_mobilenet_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..a315ccd326c2b42614b3cbde01a89a993c5af7e0
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilenet_v1/modeling_mobilenet_v1.py
@@ -0,0 +1,487 @@
+# coding=utf-8
+# Copyright 2022 Apple 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 MobileNetV1 model."""
+
+from typing import Optional, Union
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...modeling_outputs import 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_mobilenet_v1 import MobileNetV1Config
+
+
+logger = logging.get_logger(__name__)
+
+
+# General docstring
+_CONFIG_FOR_DOC = "MobileNetV1Config"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "google/mobilenet_v1_1.0_224"
+_EXPECTED_OUTPUT_SHAPE = [1, 1024, 7, 7]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "google/mobilenet_v1_1.0_224"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+def _build_tf_to_pytorch_map(model, config, tf_weights=None):
+ """
+ A map of modules from TF to PyTorch.
+ """
+
+ tf_to_pt_map = {}
+
+ if isinstance(model, MobileNetV1ForImageClassification):
+ backbone = model.mobilenet_v1
+ else:
+ backbone = model
+
+ prefix = "MobilenetV1/Conv2d_0/"
+ tf_to_pt_map[prefix + "weights"] = backbone.conv_stem.convolution.weight
+ tf_to_pt_map[prefix + "BatchNorm/beta"] = backbone.conv_stem.normalization.bias
+ tf_to_pt_map[prefix + "BatchNorm/gamma"] = backbone.conv_stem.normalization.weight
+ tf_to_pt_map[prefix + "BatchNorm/moving_mean"] = backbone.conv_stem.normalization.running_mean
+ tf_to_pt_map[prefix + "BatchNorm/moving_variance"] = backbone.conv_stem.normalization.running_var
+
+ for i in range(13):
+ tf_index = i + 1
+ pt_index = i * 2
+
+ pointer = backbone.layer[pt_index]
+ prefix = f"MobilenetV1/Conv2d_{tf_index}_depthwise/"
+ tf_to_pt_map[prefix + "depthwise_weights"] = pointer.convolution.weight
+ tf_to_pt_map[prefix + "BatchNorm/beta"] = pointer.normalization.bias
+ tf_to_pt_map[prefix + "BatchNorm/gamma"] = pointer.normalization.weight
+ tf_to_pt_map[prefix + "BatchNorm/moving_mean"] = pointer.normalization.running_mean
+ tf_to_pt_map[prefix + "BatchNorm/moving_variance"] = pointer.normalization.running_var
+
+ pointer = backbone.layer[pt_index + 1]
+ prefix = f"MobilenetV1/Conv2d_{tf_index}_pointwise/"
+ tf_to_pt_map[prefix + "weights"] = pointer.convolution.weight
+ tf_to_pt_map[prefix + "BatchNorm/beta"] = pointer.normalization.bias
+ tf_to_pt_map[prefix + "BatchNorm/gamma"] = pointer.normalization.weight
+ tf_to_pt_map[prefix + "BatchNorm/moving_mean"] = pointer.normalization.running_mean
+ tf_to_pt_map[prefix + "BatchNorm/moving_variance"] = pointer.normalization.running_var
+
+ if isinstance(model, MobileNetV1ForImageClassification):
+ prefix = "MobilenetV1/Logits/Conv2d_1c_1x1/"
+ tf_to_pt_map[prefix + "weights"] = model.classifier.weight
+ tf_to_pt_map[prefix + "biases"] = model.classifier.bias
+
+ return tf_to_pt_map
+
+
+def load_tf_weights_in_mobilenet_v1(model, config, tf_checkpoint_path):
+ """Load TensorFlow 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
+
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_checkpoint_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_checkpoint_path, name)
+ tf_weights[name] = array
+
+ # Build TF to PyTorch weights loading map
+ tf_to_pt_map = _build_tf_to_pytorch_map(model, config, tf_weights)
+
+ for name, pointer in tf_to_pt_map.items():
+ logger.info(f"Importing {name}")
+ if name not in tf_weights:
+ logger.info(f"{name} not in tf pre-trained weights, skipping")
+ continue
+
+ array = tf_weights[name]
+
+ if "depthwise_weights" in name:
+ logger.info("Transposing depthwise")
+ array = np.transpose(array, (2, 3, 0, 1))
+ elif "weights" in name:
+ logger.info("Transposing")
+ if len(pointer.shape) == 2: # copying into linear layer
+ array = array.squeeze().transpose()
+ else:
+ array = np.transpose(array, (3, 2, 0, 1))
+
+ if pointer.shape != array.shape:
+ raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
+
+ logger.info(f"Initialize PyTorch weight {name} {array.shape}")
+ pointer.data = torch.from_numpy(array)
+
+ tf_weights.pop(name, None)
+ tf_weights.pop(name + "/RMSProp", None)
+ tf_weights.pop(name + "/RMSProp_1", None)
+ tf_weights.pop(name + "/ExponentialMovingAverage", None)
+
+ logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}")
+ return model
+
+
+def apply_tf_padding(features: torch.Tensor, conv_layer: nn.Conv2d) -> torch.Tensor:
+ """
+ Apply TensorFlow-style "SAME" padding to a convolution layer. See the notes at:
+ https://www.tensorflow.org/api_docs/python/tf/nn#notes_on_padding_2
+ """
+ in_height, in_width = features.shape[-2:]
+ stride_height, stride_width = conv_layer.stride
+ kernel_height, kernel_width = conv_layer.kernel_size
+
+ if in_height % stride_height == 0:
+ pad_along_height = max(kernel_height - stride_height, 0)
+ else:
+ pad_along_height = max(kernel_height - (in_height % stride_height), 0)
+
+ if in_width % stride_width == 0:
+ pad_along_width = max(kernel_width - stride_width, 0)
+ else:
+ pad_along_width = max(kernel_width - (in_width % stride_width), 0)
+
+ pad_left = pad_along_width // 2
+ pad_right = pad_along_width - pad_left
+ pad_top = pad_along_height // 2
+ pad_bottom = pad_along_height - pad_top
+
+ padding = (pad_left, pad_right, pad_top, pad_bottom)
+ return nn.functional.pad(features, padding, "constant", 0.0)
+
+
+class MobileNetV1ConvLayer(nn.Module):
+ def __init__(
+ self,
+ config: MobileNetV1Config,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: Optional[int] = 1,
+ groups: Optional[int] = 1,
+ bias: bool = False,
+ use_normalization: Optional[bool] = True,
+ use_activation: Optional[bool or str] = True,
+ ) -> None:
+ super().__init__()
+ self.config = config
+
+ if in_channels % groups != 0:
+ raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups.")
+ if out_channels % groups != 0:
+ raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups.")
+
+ padding = 0 if config.tf_padding else int((kernel_size - 1) / 2)
+
+ self.convolution = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ groups=groups,
+ bias=bias,
+ padding_mode="zeros",
+ )
+
+ if use_normalization:
+ self.normalization = nn.BatchNorm2d(
+ num_features=out_channels,
+ eps=config.layer_norm_eps,
+ momentum=0.9997,
+ affine=True,
+ track_running_stats=True,
+ )
+ else:
+ self.normalization = None
+
+ if use_activation:
+ if isinstance(use_activation, str):
+ self.activation = ACT2FN[use_activation]
+ elif isinstance(config.hidden_act, str):
+ self.activation = ACT2FN[config.hidden_act]
+ else:
+ self.activation = config.hidden_act
+ else:
+ self.activation = None
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ if self.config.tf_padding:
+ features = apply_tf_padding(features, self.convolution)
+ features = self.convolution(features)
+ if self.normalization is not None:
+ features = self.normalization(features)
+ if self.activation is not None:
+ features = self.activation(features)
+ return features
+
+
+class MobileNetV1PreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MobileNetV1Config
+ load_tf_weights = load_tf_weights_in_mobilenet_v1
+ base_model_prefix = "mobilenet_v1"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = False
+ _no_split_modules = []
+
+ def _init_weights(self, module: Union[nn.Linear, nn.Conv2d]) -> None:
+ """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.BatchNorm2d):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+
+
+MOBILENET_V1_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 ([`MobileNetV1Config`]): 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.
+"""
+
+MOBILENET_V1_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
+ [`MobileNetV1ImageProcessor.__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 MobileNetV1 model outputting raw hidden-states without any specific head on top.",
+ MOBILENET_V1_START_DOCSTRING,
+)
+class MobileNetV1Model(MobileNetV1PreTrainedModel):
+ def __init__(self, config: MobileNetV1Config, add_pooling_layer: bool = True):
+ super().__init__(config)
+ self.config = config
+
+ depth = 32
+ out_channels = max(int(depth * config.depth_multiplier), config.min_depth)
+
+ self.conv_stem = MobileNetV1ConvLayer(
+ config,
+ in_channels=config.num_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=2,
+ )
+
+ strides = [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]
+
+ self.layer = nn.ModuleList()
+ for i in range(13):
+ in_channels = out_channels
+
+ if strides[i] == 2 or i == 0:
+ depth *= 2
+ out_channels = max(int(depth * config.depth_multiplier), config.min_depth)
+
+ self.layer.append(
+ MobileNetV1ConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=3,
+ stride=strides[i],
+ groups=in_channels,
+ )
+ )
+
+ self.layer.append(
+ MobileNetV1ConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ )
+ )
+
+ self.pooler = nn.AdaptiveAvgPool2d((1, 1)) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def _prune_heads(self, heads_to_prune):
+ raise NotImplementedError
+
+ @add_start_docstrings_to_model_forward(MOBILENET_V1_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.Tensor] = 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")
+
+ hidden_states = self.conv_stem(pixel_values)
+
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, layer_module in enumerate(self.layer):
+ hidden_states = layer_module(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ last_hidden_state = hidden_states
+
+ if self.pooler is not None:
+ pooled_output = torch.flatten(self.pooler(last_hidden_state), start_dim=1)
+ else:
+ pooled_output = None
+
+ if not return_dict:
+ return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None)
+
+ return BaseModelOutputWithPoolingAndNoAttention(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=all_hidden_states,
+ )
+
+
+@add_start_docstrings(
+ """
+ MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """,
+ MOBILENET_V1_START_DOCSTRING,
+)
+class MobileNetV1ForImageClassification(MobileNetV1PreTrainedModel):
+ def __init__(self, config: MobileNetV1Config) -> None:
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.mobilenet_v1 = MobileNetV1Model(config)
+
+ last_hidden_size = self.mobilenet_v1.layer[-1].convolution.out_channels
+
+ # Classifier head
+ self.dropout = nn.Dropout(config.classifier_dropout_prob, inplace=True)
+ self.classifier = nn.Linear(last_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(MOBILENET_V1_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,
+ output_hidden_states: Optional[bool] = None,
+ labels: Optional[torch.Tensor] = 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.mobilenet_v1(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(self.dropout(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,
+ )
+
+
+__all__ = [
+ "MobileNetV1ForImageClassification",
+ "MobileNetV1Model",
+ "MobileNetV1PreTrainedModel",
+ "load_tf_weights_in_mobilenet_v1",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__init__.py b/lib/python3.12/site-packages/transformers/models/mobilevit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..63f4f9c4720ad938a48437671f73086a3206ee11
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_mobilevit import *
+ from .feature_extraction_mobilevit import *
+ from .image_processing_mobilevit import *
+ from .modeling_mobilevit import *
+ from .modeling_tf_mobilevit import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..afd3183300e2cde10417fb157686c34baf86f10b
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/configuration_mobilevit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/configuration_mobilevit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2405a1313f81508fe930c0b8cd742f6482a6d086
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/configuration_mobilevit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/feature_extraction_mobilevit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/feature_extraction_mobilevit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5c76998b97c9aac9f84ff016b7777307a1999ece
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/feature_extraction_mobilevit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/image_processing_mobilevit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/image_processing_mobilevit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ea4a903268679bd039b5c05613b7255b94efe3e7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/image_processing_mobilevit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_mobilevit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_mobilevit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..435915d530cff9b0ac8e7bb294974e5ea8b93646
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_mobilevit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_tf_mobilevit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_tf_mobilevit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..08e6649fa1fe6df85afa97f939a75a41d1ffd2d7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mobilevit/__pycache__/modeling_tf_mobilevit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/configuration_mobilevit.py b/lib/python3.12/site-packages/transformers/models/mobilevit/configuration_mobilevit.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ef891a96f6ccce3af06f13dba6644b18f291421
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/configuration_mobilevit.py
@@ -0,0 +1,172 @@
+# 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.
+"""MobileViT 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__)
+
+
+class MobileViTConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MobileViTModel`]. It is used to instantiate a
+ MobileViT 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 MobileViT
+ [apple/mobilevit-small](https://huggingface.co/apple/mobilevit-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:
+ num_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ image_size (`int`, *optional*, defaults to 256):
+ The size (resolution) of each image.
+ patch_size (`int`, *optional*, defaults to 2):
+ The size (resolution) of each patch.
+ hidden_sizes (`List[int]`, *optional*, defaults to `[144, 192, 240]`):
+ Dimensionality (hidden size) of the Transformer encoders at each stage.
+ neck_hidden_sizes (`List[int]`, *optional*, defaults to `[16, 32, 64, 96, 128, 160, 640]`):
+ The number of channels for the feature maps of the backbone.
+ num_attention_heads (`int`, *optional*, defaults to 4):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ mlp_ratio (`float`, *optional*, defaults to 2.0):
+ The ratio of the number of channels in the output of the MLP to the number of channels in the input.
+ expand_ratio (`float`, *optional*, defaults to 4.0):
+ Expansion factor for the MobileNetv2 layers.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the Transformer encoder and convolution layers.
+ conv_kernel_size (`int`, *optional*, defaults to 3):
+ The size of the convolutional kernel in the MobileViT layer.
+ output_stride (`int`, *optional*, defaults to 32):
+ The ratio of the spatial resolution of the output to the resolution of the input image.
+ hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
+ The dropout probability for all fully connected layers in the Transformer encoder.
+ attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ classifier_dropout_prob (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for attached classifiers.
+ 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.
+ qkv_bias (`bool`, *optional*, defaults to `True`):
+ Whether to add a bias to the queries, keys and values.
+ aspp_out_channels (`int`, *optional*, defaults to 256):
+ Number of output channels used in the ASPP layer for semantic segmentation.
+ atrous_rates (`List[int]`, *optional*, defaults to `[6, 12, 18]`):
+ Dilation (atrous) factors used in the ASPP layer for semantic segmentation.
+ aspp_dropout_prob (`float`, *optional*, defaults to 0.1):
+ The dropout ratio for the ASPP layer for semantic segmentation.
+ 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 MobileViTConfig, MobileViTModel
+
+ >>> # Initializing a mobilevit-small style configuration
+ >>> configuration = MobileViTConfig()
+
+ >>> # Initializing a model from the mobilevit-small style configuration
+ >>> model = MobileViTModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "mobilevit"
+
+ def __init__(
+ self,
+ num_channels=3,
+ image_size=256,
+ patch_size=2,
+ hidden_sizes=[144, 192, 240],
+ neck_hidden_sizes=[16, 32, 64, 96, 128, 160, 640],
+ num_attention_heads=4,
+ mlp_ratio=2.0,
+ expand_ratio=4.0,
+ hidden_act="silu",
+ conv_kernel_size=3,
+ output_stride=32,
+ hidden_dropout_prob=0.1,
+ attention_probs_dropout_prob=0.0,
+ classifier_dropout_prob=0.1,
+ initializer_range=0.02,
+ layer_norm_eps=1e-5,
+ qkv_bias=True,
+ aspp_out_channels=256,
+ atrous_rates=[6, 12, 18],
+ aspp_dropout_prob=0.1,
+ semantic_loss_ignore_index=255,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.num_channels = num_channels
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.hidden_sizes = hidden_sizes
+ self.neck_hidden_sizes = neck_hidden_sizes
+ self.num_attention_heads = num_attention_heads
+ self.mlp_ratio = mlp_ratio
+ self.expand_ratio = expand_ratio
+ self.hidden_act = hidden_act
+ self.conv_kernel_size = conv_kernel_size
+ self.output_stride = output_stride
+ self.hidden_dropout_prob = hidden_dropout_prob
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
+ self.classifier_dropout_prob = classifier_dropout_prob
+ self.initializer_range = initializer_range
+ self.layer_norm_eps = layer_norm_eps
+ self.qkv_bias = qkv_bias
+
+ # decode head attributes for semantic segmentation
+ self.aspp_out_channels = aspp_out_channels
+ self.atrous_rates = atrous_rates
+ self.aspp_dropout_prob = aspp_dropout_prob
+ self.semantic_loss_ignore_index = semantic_loss_ignore_index
+
+
+class MobileViTOnnxConfig(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 outputs(self) -> Mapping[str, Mapping[int, str]]:
+ if self.task == "image-classification":
+ return OrderedDict([("logits", {0: "batch"})])
+ else:
+ return OrderedDict([("last_hidden_state", {0: "batch"}), ("pooler_output", {0: "batch"})])
+
+ @property
+ def atol_for_validation(self) -> float:
+ return 1e-4
+
+
+__all__ = ["MobileViTConfig", "MobileViTOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/feature_extraction_mobilevit.py b/lib/python3.12/site-packages/transformers/models/mobilevit/feature_extraction_mobilevit.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb98b2f6e145227fa60a7793cd6770cd1d9434b6
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/feature_extraction_mobilevit.py
@@ -0,0 +1,36 @@
+# 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 MobileViT."""
+
+import warnings
+
+from ...utils import logging
+from .image_processing_mobilevit import MobileViTImageProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+class MobileViTFeatureExtractor(MobileViTImageProcessor):
+ def __init__(self, *args, **kwargs) -> None:
+ warnings.warn(
+ "The class MobileViTFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
+ " Please use MobileViTImageProcessor instead.",
+ FutureWarning,
+ )
+ super().__init__(*args, **kwargs)
+
+
+__all__ = ["MobileViTFeatureExtractor"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/image_processing_mobilevit.py b/lib/python3.12/site-packages/transformers/models/mobilevit/image_processing_mobilevit.py
new file mode 100644
index 0000000000000000000000000000000000000000..f59c246627c3b48bc2c9d68bc8218f952f548d2b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/image_processing_mobilevit.py
@@ -0,0 +1,485 @@
+# 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 MobileViT."""
+
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from ...image_transforms import flip_channel_order, get_resize_output_image_size, resize, to_channel_dimension_format
+from ...image_utils import (
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ infer_channel_dimension_format,
+ is_scaled_image,
+ make_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import (
+ TensorType,
+ filter_out_non_signature_kwargs,
+ is_torch_available,
+ is_torch_tensor,
+ is_vision_available,
+ logging,
+)
+
+
+if is_vision_available():
+ import PIL
+
+if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class MobileViTImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a MobileViT 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 the `preprocess` method.
+ size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 224}`):
+ Controls the size of the output image after resizing. Can be overridden by the `size` parameter in the
+ `preprocess` method.
+ resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
+ Defines the resampling filter to use if resizing the image. Can be overridden by the `resample` 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.
+ 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_center_crop (`bool`, *optional*, defaults to `True`):
+ Whether to crop the input at the center. 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 the `do_center_crop` parameter in
+ the `preprocess` method.
+ crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 256, "width": 256}`):
+ Desired output size `(size["height"], size["width"])` when applying center-cropping. Can be overridden by
+ the `crop_size` parameter in the `preprocess` method.
+ do_flip_channel_order (`bool`, *optional*, defaults to `True`):
+ Whether to flip the color channels from RGB to BGR. Can be overridden by the `do_flip_channel_order`
+ parameter in the `preprocess` method.
+ """
+
+ model_input_names = ["pixel_values"]
+
+ def __init__(
+ self,
+ 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_center_crop: bool = True,
+ crop_size: Dict[str, int] = None,
+ do_flip_channel_order: bool = True,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ size = size if size is not None else {"shortest_edge": 224}
+ size = get_size_dict(size, default_to_square=False)
+ crop_size = crop_size if crop_size is not None else {"height": 256, "width": 256}
+ crop_size = get_size_dict(crop_size, param_name="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_flip_channel_order = do_flip_channel_order
+
+ # Copied from transformers.models.mobilenet_v1.image_processing_mobilenet_v1.MobileNetV1ImageProcessor.resize with PILImageResampling.BICUBIC->PILImageResampling.BILINEAR
+ 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. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
+ resized to keep the input aspect ratio.
+
+ Args:
+ image (`np.ndarray`):
+ Image to resize.
+ size (`Dict[str, int]`):
+ Size of the output image.
+ resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
+ 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.
+ """
+ default_to_square = True
+ if "shortest_edge" in size:
+ size = size["shortest_edge"]
+ default_to_square = False
+ elif "height" in size and "width" in size:
+ size = (size["height"], size["width"])
+ else:
+ raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
+
+ output_size = get_resize_output_image_size(
+ image,
+ size=size,
+ default_to_square=default_to_square,
+ 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 flip_channel_order(
+ self,
+ image: np.ndarray,
+ data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> np.ndarray:
+ """
+ Flip the color channels from RGB to BGR or vice versa.
+
+ Args:
+ image (`np.ndarray`):
+ The image, represented as a numpy array.
+ data_format (`ChannelDimension` or `str`, *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.
+ """
+ return flip_channel_order(image, data_format=data_format, input_data_format=input_data_format)
+
+ def __call__(self, images, segmentation_maps=None, **kwargs):
+ """
+ Preprocesses a batch of images and optionally segmentation maps.
+
+ Overrides the `__call__` method of the `Preprocessor` class so that both images and segmentation maps can be
+ passed in as positional arguments.
+ """
+ return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)
+
+ def _preprocess(
+ self,
+ image: ImageInput,
+ do_resize: bool,
+ do_rescale: bool,
+ do_center_crop: bool,
+ do_flip_channel_order: bool,
+ size: Optional[Dict[str, int]] = None,
+ resample: PILImageResampling = None,
+ rescale_factor: Optional[float] = None,
+ crop_size: Optional[Dict[str, int]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ):
+ if do_resize:
+ image = self.resize(image=image, size=size, resample=resample, 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_center_crop:
+ image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)
+
+ if do_flip_channel_order:
+ image = self.flip_channel_order(image, input_data_format=input_data_format)
+
+ return image
+
+ def _preprocess_image(
+ self,
+ image: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = None,
+ do_rescale: Optional[bool] = None,
+ rescale_factor: Optional[float] = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: Dict[str, int] = None,
+ do_flip_channel_order: Optional[bool] = None,
+ data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> np.ndarray:
+ """Preprocesses a single image."""
+ # All transformations expect numpy arrays.
+ image = to_numpy_array(image)
+ if do_rescale and is_scaled_image(image):
+ 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:
+ input_data_format = infer_channel_dimension_format(image)
+
+ image = self._preprocess(
+ image=image,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_flip_channel_order=do_flip_channel_order,
+ input_data_format=input_data_format,
+ )
+
+ image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+
+ return image
+
+ def _preprocess_mask(
+ self,
+ segmentation_map: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: Dict[str, int] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> np.ndarray:
+ """Preprocesses a single mask."""
+ segmentation_map = to_numpy_array(segmentation_map)
+ # Add channel dimension if missing - needed for certain transformations
+ if segmentation_map.ndim == 2:
+ added_channel_dim = True
+ segmentation_map = segmentation_map[None, ...]
+ input_data_format = ChannelDimension.FIRST
+ else:
+ added_channel_dim = False
+ if input_data_format is None:
+ input_data_format = infer_channel_dimension_format(segmentation_map, num_channels=1)
+
+ segmentation_map = self._preprocess(
+ image=segmentation_map,
+ do_resize=do_resize,
+ size=size,
+ resample=PILImageResampling.NEAREST,
+ do_rescale=False,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_flip_channel_order=False,
+ input_data_format=input_data_format,
+ )
+ # Remove extra channel dimension if added for processing
+ if added_channel_dim:
+ segmentation_map = segmentation_map.squeeze(0)
+ segmentation_map = segmentation_map.astype(np.int64)
+ return segmentation_map
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ segmentation_maps: Optional[ImageInput] = None,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = None,
+ do_rescale: Optional[bool] = None,
+ rescale_factor: Optional[float] = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: Dict[str, int] = None,
+ do_flip_channel_order: Optional[bool] = None,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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`.
+ segmentation_maps (`ImageInput`, *optional*):
+ Segmentation map to preprocess.
+ 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 (`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 by rescale factor.
+ rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+ Rescale factor to rescale the image by if `do_rescale` 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 if `do_center_crop` is set to `True`.
+ do_flip_channel_order (`bool`, *optional*, defaults to `self.do_flip_channel_order`):
+ Whether to flip the channel order of the image.
+ 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
+ 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_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
+ do_flip_channel_order = (
+ do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
+ )
+
+ size = size if size is not None else self.size
+ size = get_size_dict(size, default_to_square=False)
+ 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)
+
+ if segmentation_maps is not None:
+ segmentation_maps = make_list_of_images(segmentation_maps, expected_ndims=2)
+
+ 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."
+ )
+
+ if segmentation_maps is not None and not valid_images(segmentation_maps):
+ raise ValueError(
+ "Invalid segmentation map 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_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ )
+
+ images = [
+ self._preprocess_image(
+ image=img,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_flip_channel_order=do_flip_channel_order,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+ for img in images
+ ]
+
+ data = {"pixel_values": images}
+
+ if segmentation_maps is not None:
+ segmentation_maps = [
+ self._preprocess_mask(
+ segmentation_map=segmentation_map,
+ do_resize=do_resize,
+ size=size,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ input_data_format=input_data_format,
+ )
+ for segmentation_map in segmentation_maps
+ ]
+
+ data["labels"] = segmentation_maps
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+ # Copied from transformers.models.beit.image_processing_beit.BeitImageProcessor.post_process_semantic_segmentation with Beit->MobileViT
+ def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
+ """
+ Converts the output of [`MobileViTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.
+
+ Args:
+ outputs ([`MobileViTForSemanticSegmentation`]):
+ 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
+
+
+__all__ = ["MobileViTImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_mobilevit.py b/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_mobilevit.py
new file mode 100644
index 0000000000000000000000000000000000000000..f41da2bafaf280f5a5596e23344489f78c70f6cc
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_mobilevit.py
@@ -0,0 +1,1080 @@
+# coding=utf-8
+# Copyright 2022 Apple 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.
+#
+# Original license: https://github.com/apple/ml-cvnets/blob/main/LICENSE
+"""PyTorch MobileViT model."""
+
+import math
+from typing import Dict, 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 (
+ BaseModelOutputWithNoAttention,
+ BaseModelOutputWithPoolingAndNoAttention,
+ ImageClassifierOutputWithNoAttention,
+ SemanticSegmenterOutput,
+)
+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,
+ torch_int,
+)
+from .configuration_mobilevit import MobileViTConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# General docstring
+_CONFIG_FOR_DOC = "MobileViTConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "apple/mobilevit-small"
+_EXPECTED_OUTPUT_SHAPE = [1, 640, 8, 8]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "apple/mobilevit-small"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+def make_divisible(value: int, divisor: int = 8, min_value: Optional[int] = None) -> int:
+ """
+ Ensure that all layers have a channel count that is divisible by `divisor`. This function is taken from the
+ original TensorFlow repo. It can be seen here:
+ https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+ """
+ if min_value is None:
+ min_value = divisor
+ new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
+ # Make sure that round down does not go down by more than 10%.
+ if new_value < 0.9 * value:
+ new_value += divisor
+ return int(new_value)
+
+
+class MobileViTConvLayer(nn.Module):
+ def __init__(
+ self,
+ config: MobileViTConfig,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ groups: int = 1,
+ bias: bool = False,
+ dilation: int = 1,
+ use_normalization: bool = True,
+ use_activation: Union[bool, str] = True,
+ ) -> None:
+ super().__init__()
+ padding = int((kernel_size - 1) / 2) * dilation
+
+ if in_channels % groups != 0:
+ raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups.")
+ if out_channels % groups != 0:
+ raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups.")
+
+ self.convolution = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode="zeros",
+ )
+
+ if use_normalization:
+ self.normalization = nn.BatchNorm2d(
+ num_features=out_channels,
+ eps=1e-5,
+ momentum=0.1,
+ affine=True,
+ track_running_stats=True,
+ )
+ else:
+ self.normalization = None
+
+ if use_activation:
+ if isinstance(use_activation, str):
+ self.activation = ACT2FN[use_activation]
+ elif isinstance(config.hidden_act, str):
+ self.activation = ACT2FN[config.hidden_act]
+ else:
+ self.activation = config.hidden_act
+ else:
+ self.activation = None
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ features = self.convolution(features)
+ if self.normalization is not None:
+ features = self.normalization(features)
+ if self.activation is not None:
+ features = self.activation(features)
+ return features
+
+
+class MobileViTInvertedResidual(nn.Module):
+ """
+ Inverted residual block (MobileNetv2): https://arxiv.org/abs/1801.04381
+ """
+
+ def __init__(
+ self, config: MobileViTConfig, in_channels: int, out_channels: int, stride: int, dilation: int = 1
+ ) -> None:
+ super().__init__()
+ expanded_channels = make_divisible(int(round(in_channels * config.expand_ratio)), 8)
+
+ if stride not in [1, 2]:
+ raise ValueError(f"Invalid stride {stride}.")
+
+ self.use_residual = (stride == 1) and (in_channels == out_channels)
+
+ self.expand_1x1 = MobileViTConvLayer(
+ config, in_channels=in_channels, out_channels=expanded_channels, kernel_size=1
+ )
+
+ self.conv_3x3 = MobileViTConvLayer(
+ config,
+ in_channels=expanded_channels,
+ out_channels=expanded_channels,
+ kernel_size=3,
+ stride=stride,
+ groups=expanded_channels,
+ dilation=dilation,
+ )
+
+ self.reduce_1x1 = MobileViTConvLayer(
+ config,
+ in_channels=expanded_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ use_activation=False,
+ )
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ residual = features
+
+ features = self.expand_1x1(features)
+ features = self.conv_3x3(features)
+ features = self.reduce_1x1(features)
+
+ return residual + features if self.use_residual else features
+
+
+class MobileViTMobileNetLayer(nn.Module):
+ def __init__(
+ self, config: MobileViTConfig, in_channels: int, out_channels: int, stride: int = 1, num_stages: int = 1
+ ) -> None:
+ super().__init__()
+
+ self.layer = nn.ModuleList()
+ for i in range(num_stages):
+ layer = MobileViTInvertedResidual(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=stride if i == 0 else 1,
+ )
+ self.layer.append(layer)
+ in_channels = out_channels
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ for layer_module in self.layer:
+ features = layer_module(features)
+ return features
+
+
+class MobileViTSelfAttention(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int) -> None:
+ super().__init__()
+
+ if hidden_size % config.num_attention_heads != 0:
+ raise ValueError(
+ f"The hidden size {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(hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias)
+ self.key = nn.Linear(hidden_size, self.all_head_size, bias=config.qkv_bias)
+ self.value = nn.Linear(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) -> 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)
+
+ 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)
+ return context_layer
+
+
+class MobileViTSelfOutput(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int) -> None:
+ super().__init__()
+ self.dense = nn.Linear(hidden_size, 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
+
+
+class MobileViTAttention(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int) -> None:
+ super().__init__()
+ self.attention = MobileViTSelfAttention(config, hidden_size)
+ self.output = MobileViTSelfOutput(config, hidden_size)
+ 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) -> torch.Tensor:
+ self_outputs = self.attention(hidden_states)
+ attention_output = self.output(self_outputs)
+ return attention_output
+
+
+class MobileViTIntermediate(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None:
+ super().__init__()
+ self.dense = nn.Linear(hidden_size, 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 MobileViTOutput(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None:
+ super().__init__()
+ self.dense = nn.Linear(intermediate_size, 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
+
+
+class MobileViTTransformerLayer(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int) -> None:
+ super().__init__()
+ self.attention = MobileViTAttention(config, hidden_size)
+ self.intermediate = MobileViTIntermediate(config, hidden_size, intermediate_size)
+ self.output = MobileViTOutput(config, hidden_size, intermediate_size)
+ self.layernorm_before = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+ self.layernorm_after = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ attention_output = self.attention(self.layernorm_before(hidden_states))
+ hidden_states = attention_output + hidden_states
+
+ layer_output = self.layernorm_after(hidden_states)
+ layer_output = self.intermediate(layer_output)
+ layer_output = self.output(layer_output, hidden_states)
+ return layer_output
+
+
+class MobileViTTransformer(nn.Module):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, num_stages: int) -> None:
+ super().__init__()
+
+ self.layer = nn.ModuleList()
+ for _ in range(num_stages):
+ transformer_layer = MobileViTTransformerLayer(
+ config,
+ hidden_size=hidden_size,
+ intermediate_size=int(hidden_size * config.mlp_ratio),
+ )
+ self.layer.append(transformer_layer)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for layer_module in self.layer:
+ hidden_states = layer_module(hidden_states)
+ return hidden_states
+
+
+class MobileViTLayer(nn.Module):
+ """
+ MobileViT block: https://arxiv.org/abs/2110.02178
+ """
+
+ def __init__(
+ self,
+ config: MobileViTConfig,
+ in_channels: int,
+ out_channels: int,
+ stride: int,
+ hidden_size: int,
+ num_stages: int,
+ dilation: int = 1,
+ ) -> None:
+ super().__init__()
+ self.patch_width = config.patch_size
+ self.patch_height = config.patch_size
+
+ if stride == 2:
+ self.downsampling_layer = MobileViTInvertedResidual(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=stride if dilation == 1 else 1,
+ dilation=dilation // 2 if dilation > 1 else 1,
+ )
+ in_channels = out_channels
+ else:
+ self.downsampling_layer = None
+
+ self.conv_kxk = MobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=config.conv_kernel_size,
+ )
+
+ self.conv_1x1 = MobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=hidden_size,
+ kernel_size=1,
+ use_normalization=False,
+ use_activation=False,
+ )
+
+ self.transformer = MobileViTTransformer(
+ config,
+ hidden_size=hidden_size,
+ num_stages=num_stages,
+ )
+
+ self.layernorm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+
+ self.conv_projection = MobileViTConvLayer(
+ config, in_channels=hidden_size, out_channels=in_channels, kernel_size=1
+ )
+
+ self.fusion = MobileViTConvLayer(
+ config, in_channels=2 * in_channels, out_channels=in_channels, kernel_size=config.conv_kernel_size
+ )
+
+ def unfolding(self, features: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
+ patch_width, patch_height = self.patch_width, self.patch_height
+ patch_area = int(patch_width * patch_height)
+
+ batch_size, channels, orig_height, orig_width = features.shape
+
+ new_height = (
+ torch_int(torch.ceil(orig_height / patch_height) * patch_height)
+ if torch.jit.is_tracing()
+ else int(math.ceil(orig_height / patch_height) * patch_height)
+ )
+ new_width = (
+ torch_int(torch.ceil(orig_width / patch_width) * patch_width)
+ if torch.jit.is_tracing()
+ else int(math.ceil(orig_width / patch_width) * patch_width)
+ )
+
+ interpolate = False
+ if new_width != orig_width or new_height != orig_height:
+ # Note: Padding can be done, but then it needs to be handled in attention function.
+ features = nn.functional.interpolate(
+ features, size=(new_height, new_width), mode="bilinear", align_corners=False
+ )
+ interpolate = True
+
+ # number of patches along width and height
+ num_patch_width = new_width // patch_width
+ num_patch_height = new_height // patch_height
+ num_patches = num_patch_height * num_patch_width
+
+ # convert from shape (batch_size, channels, orig_height, orig_width)
+ # to the shape (batch_size * patch_area, num_patches, channels)
+ patches = features.reshape(
+ batch_size * channels * num_patch_height, patch_height, num_patch_width, patch_width
+ )
+ patches = patches.transpose(1, 2)
+ patches = patches.reshape(batch_size, channels, num_patches, patch_area)
+ patches = patches.transpose(1, 3)
+ patches = patches.reshape(batch_size * patch_area, num_patches, -1)
+
+ info_dict = {
+ "orig_size": (orig_height, orig_width),
+ "batch_size": batch_size,
+ "channels": channels,
+ "interpolate": interpolate,
+ "num_patches": num_patches,
+ "num_patches_width": num_patch_width,
+ "num_patches_height": num_patch_height,
+ }
+ return patches, info_dict
+
+ def folding(self, patches: torch.Tensor, info_dict: Dict) -> torch.Tensor:
+ patch_width, patch_height = self.patch_width, self.patch_height
+ patch_area = int(patch_width * patch_height)
+
+ batch_size = info_dict["batch_size"]
+ channels = info_dict["channels"]
+ num_patches = info_dict["num_patches"]
+ num_patch_height = info_dict["num_patches_height"]
+ num_patch_width = info_dict["num_patches_width"]
+
+ # convert from shape (batch_size * patch_area, num_patches, channels)
+ # back to shape (batch_size, channels, orig_height, orig_width)
+ features = patches.contiguous().view(batch_size, patch_area, num_patches, -1)
+ features = features.transpose(1, 3)
+ features = features.reshape(
+ batch_size * channels * num_patch_height, num_patch_width, patch_height, patch_width
+ )
+ features = features.transpose(1, 2)
+ features = features.reshape(
+ batch_size, channels, num_patch_height * patch_height, num_patch_width * patch_width
+ )
+
+ if info_dict["interpolate"]:
+ features = nn.functional.interpolate(
+ features, size=info_dict["orig_size"], mode="bilinear", align_corners=False
+ )
+
+ return features
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ # reduce spatial dimensions if needed
+ if self.downsampling_layer:
+ features = self.downsampling_layer(features)
+
+ residual = features
+
+ # local representation
+ features = self.conv_kxk(features)
+ features = self.conv_1x1(features)
+
+ # convert feature map to patches
+ patches, info_dict = self.unfolding(features)
+
+ # learn global representations
+ patches = self.transformer(patches)
+ patches = self.layernorm(patches)
+
+ # convert patches back to feature maps
+ features = self.folding(patches, info_dict)
+
+ features = self.conv_projection(features)
+ features = self.fusion(torch.cat((residual, features), dim=1))
+ return features
+
+
+class MobileViTEncoder(nn.Module):
+ def __init__(self, config: MobileViTConfig) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layer = nn.ModuleList()
+ self.gradient_checkpointing = False
+
+ # segmentation architectures like DeepLab and PSPNet modify the strides
+ # of the classification backbones
+ dilate_layer_4 = dilate_layer_5 = False
+ if config.output_stride == 8:
+ dilate_layer_4 = True
+ dilate_layer_5 = True
+ elif config.output_stride == 16:
+ dilate_layer_5 = True
+
+ dilation = 1
+
+ layer_1 = MobileViTMobileNetLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[0],
+ out_channels=config.neck_hidden_sizes[1],
+ stride=1,
+ num_stages=1,
+ )
+ self.layer.append(layer_1)
+
+ layer_2 = MobileViTMobileNetLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[1],
+ out_channels=config.neck_hidden_sizes[2],
+ stride=2,
+ num_stages=3,
+ )
+ self.layer.append(layer_2)
+
+ layer_3 = MobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[2],
+ out_channels=config.neck_hidden_sizes[3],
+ stride=2,
+ hidden_size=config.hidden_sizes[0],
+ num_stages=2,
+ )
+ self.layer.append(layer_3)
+
+ if dilate_layer_4:
+ dilation *= 2
+
+ layer_4 = MobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[3],
+ out_channels=config.neck_hidden_sizes[4],
+ stride=2,
+ hidden_size=config.hidden_sizes[1],
+ num_stages=4,
+ dilation=dilation,
+ )
+ self.layer.append(layer_4)
+
+ if dilate_layer_5:
+ dilation *= 2
+
+ layer_5 = MobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[4],
+ out_channels=config.neck_hidden_sizes[5],
+ stride=2,
+ hidden_size=config.hidden_sizes[2],
+ num_stages=3,
+ dilation=dilation,
+ )
+ self.layer.append(layer_5)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ ) -> Union[tuple, BaseModelOutputWithNoAttention]:
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, layer_module in enumerate(self.layer):
+ if self.gradient_checkpointing and self.training:
+ hidden_states = self._gradient_checkpointing_func(
+ layer_module.__call__,
+ hidden_states,
+ )
+ else:
+ 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 MobileViTPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MobileViTConfig
+ base_model_prefix = "mobilevit"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MobileViTLayer"]
+
+ 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.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+
+
+MOBILEVIT_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 ([`MobileViTConfig`]): 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.
+"""
+
+MOBILEVIT_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
+ [`MobileViTImageProcessor.__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 MobileViT model outputting raw hidden-states without any specific head on top.",
+ MOBILEVIT_START_DOCSTRING,
+)
+class MobileViTModel(MobileViTPreTrainedModel):
+ def __init__(self, config: MobileViTConfig, expand_output: bool = True):
+ super().__init__(config)
+ self.config = config
+ self.expand_output = expand_output
+
+ self.conv_stem = MobileViTConvLayer(
+ config,
+ in_channels=config.num_channels,
+ out_channels=config.neck_hidden_sizes[0],
+ kernel_size=3,
+ stride=2,
+ )
+
+ self.encoder = MobileViTEncoder(config)
+
+ if self.expand_output:
+ self.conv_1x1_exp = MobileViTConvLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[5],
+ out_channels=config.neck_hidden_sizes[6],
+ kernel_size=1,
+ )
+
+ # 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_index, heads in heads_to_prune.items():
+ mobilevit_layer = self.encoder.layer[layer_index]
+ if isinstance(mobilevit_layer, MobileViTLayer):
+ for transformer_layer in mobilevit_layer.transformer.layer:
+ transformer_layer.attention.prune_heads(heads)
+
+ @add_start_docstrings_to_model_forward(MOBILEVIT_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.Tensor] = 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.conv_stem(pixel_values)
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ if self.expand_output:
+ last_hidden_state = self.conv_1x1_exp(encoder_outputs[0])
+
+ # global average pooling: (batch_size, channels, height, width) -> (batch_size, channels)
+ pooled_output = torch.mean(last_hidden_state, dim=[-2, -1], keepdim=False)
+ else:
+ last_hidden_state = encoder_outputs[0]
+ pooled_output = None
+
+ if not return_dict:
+ output = (last_hidden_state, pooled_output) if pooled_output is not None else (last_hidden_state,)
+ return 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(
+ """
+ MobileViT model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """,
+ MOBILEVIT_START_DOCSTRING,
+)
+class MobileViTForImageClassification(MobileViTPreTrainedModel):
+ def __init__(self, config: MobileViTConfig) -> None:
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.mobilevit = MobileViTModel(config)
+
+ # Classifier head
+ self.dropout = nn.Dropout(config.classifier_dropout_prob, inplace=True)
+ self.classifier = (
+ nn.Linear(config.neck_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(MOBILEVIT_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,
+ output_hidden_states: Optional[bool] = None,
+ labels: Optional[torch.Tensor] = 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.mobilevit(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(self.dropout(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,
+ )
+
+
+class MobileViTASPPPooling(nn.Module):
+ def __init__(self, config: MobileViTConfig, in_channels: int, out_channels: int) -> None:
+ super().__init__()
+
+ self.global_pool = nn.AdaptiveAvgPool2d(output_size=1)
+
+ self.conv_1x1 = MobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ stride=1,
+ use_normalization=True,
+ use_activation="relu",
+ )
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ spatial_size = features.shape[-2:]
+ features = self.global_pool(features)
+ features = self.conv_1x1(features)
+ features = nn.functional.interpolate(features, size=spatial_size, mode="bilinear", align_corners=False)
+ return features
+
+
+class MobileViTASPP(nn.Module):
+ """
+ ASPP module defined in DeepLab papers: https://arxiv.org/abs/1606.00915, https://arxiv.org/abs/1706.05587
+ """
+
+ def __init__(self, config: MobileViTConfig) -> None:
+ super().__init__()
+
+ in_channels = config.neck_hidden_sizes[-2]
+ out_channels = config.aspp_out_channels
+
+ if len(config.atrous_rates) != 3:
+ raise ValueError("Expected 3 values for atrous_rates")
+
+ self.convs = nn.ModuleList()
+
+ in_projection = MobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ use_activation="relu",
+ )
+ self.convs.append(in_projection)
+
+ self.convs.extend(
+ [
+ MobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ dilation=rate,
+ use_activation="relu",
+ )
+ for rate in config.atrous_rates
+ ]
+ )
+
+ pool_layer = MobileViTASPPPooling(config, in_channels, out_channels)
+ self.convs.append(pool_layer)
+
+ self.project = MobileViTConvLayer(
+ config, in_channels=5 * out_channels, out_channels=out_channels, kernel_size=1, use_activation="relu"
+ )
+
+ self.dropout = nn.Dropout(p=config.aspp_dropout_prob)
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ pyramid = []
+ for conv in self.convs:
+ pyramid.append(conv(features))
+ pyramid = torch.cat(pyramid, dim=1)
+
+ pooled_features = self.project(pyramid)
+ pooled_features = self.dropout(pooled_features)
+ return pooled_features
+
+
+class MobileViTDeepLabV3(nn.Module):
+ """
+ DeepLabv3 architecture: https://arxiv.org/abs/1706.05587
+ """
+
+ def __init__(self, config: MobileViTConfig) -> None:
+ super().__init__()
+ self.aspp = MobileViTASPP(config)
+
+ self.dropout = nn.Dropout2d(config.classifier_dropout_prob)
+
+ self.classifier = MobileViTConvLayer(
+ config,
+ in_channels=config.aspp_out_channels,
+ out_channels=config.num_labels,
+ kernel_size=1,
+ use_normalization=False,
+ use_activation=False,
+ bias=True,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ features = self.aspp(hidden_states[-1])
+ features = self.dropout(features)
+ features = self.classifier(features)
+ return features
+
+
+@add_start_docstrings(
+ """
+ MobileViT model with a semantic segmentation head on top, e.g. for Pascal VOC.
+ """,
+ MOBILEVIT_START_DOCSTRING,
+)
+class MobileViTForSemanticSegmentation(MobileViTPreTrainedModel):
+ def __init__(self, config: MobileViTConfig) -> None:
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.mobilevit = MobileViTModel(config, expand_output=False)
+ self.segmentation_head = MobileViTDeepLabV3(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(MOBILEVIT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=SemanticSegmenterOutput, config_class=_CONFIG_FOR_DOC)
+ 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, 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
+ >>> import requests
+ >>> import torch
+ >>> from PIL import Image
+ >>> from transformers import AutoImageProcessor, MobileViTForSemanticSegmentation
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-small")
+ >>> model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-small")
+
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> with torch.no_grad():
+ ... outputs = model(**inputs)
+
+ >>> # logits are of shape (batch_size, num_labels, height, width)
+ >>> logits = outputs.logits
+ ```"""
+ 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 labels is not None and self.config.num_labels == 1:
+ raise ValueError("The number of labels should be greater than one")
+
+ outputs = self.mobilevit(
+ pixel_values,
+ 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]
+
+ logits = self.segmentation_head(encoder_hidden_states)
+
+ loss = None
+ if labels is not None:
+ # upsample logits to the images' original size
+ upsampled_logits = nn.functional.interpolate(
+ logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
+ )
+ loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index)
+ loss = loss_fct(upsampled_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=None,
+ )
+
+
+__all__ = [
+ "MobileViTForImageClassification",
+ "MobileViTForSemanticSegmentation",
+ "MobileViTModel",
+ "MobileViTPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_tf_mobilevit.py b/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_tf_mobilevit.py
new file mode 100644
index 0000000000000000000000000000000000000000..76397f160b59a02348473a97e0577244481f461b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mobilevit/modeling_tf_mobilevit.py
@@ -0,0 +1,1378 @@
+# coding=utf-8
+# Copyright 2022 Apple 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.
+#
+# Original license: https://github.com/apple/ml-cvnets/blob/main/LICENSE
+"""TensorFlow 2.0 MobileViT model."""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple, Union
+
+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,
+ replace_return_docstrings,
+)
+from ...modeling_tf_outputs import (
+ TFBaseModelOutput,
+ TFBaseModelOutputWithPooling,
+ TFImageClassifierOutputWithNoAttention,
+ TFSemanticSegmenterOutputWithNoAttention,
+)
+from ...modeling_tf_utils import (
+ TFPreTrainedModel,
+ TFSequenceClassificationLoss,
+ keras,
+ keras_serializable,
+ unpack_inputs,
+)
+from ...tf_utils import shape_list, stable_softmax
+from ...utils import logging
+from .configuration_mobilevit import MobileViTConfig
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "MobileViTConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "apple/mobilevit-small"
+_EXPECTED_OUTPUT_SHAPE = [1, 640, 8, 8]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "apple/mobilevit-small"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+def make_divisible(value: int, divisor: int = 8, min_value: Optional[int] = None) -> int:
+ """
+ Ensure that all layers have a channel count that is divisible by `divisor`. This function is taken from the
+ original TensorFlow repo. It can be seen here:
+ https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+ """
+ if min_value is None:
+ min_value = divisor
+ new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
+ # Make sure that round down does not go down by more than 10%.
+ if new_value < 0.9 * value:
+ new_value += divisor
+ return int(new_value)
+
+
+class TFMobileViTConvLayer(keras.layers.Layer):
+ def __init__(
+ self,
+ config: MobileViTConfig,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ groups: int = 1,
+ bias: bool = False,
+ dilation: int = 1,
+ use_normalization: bool = True,
+ use_activation: Union[bool, str] = True,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ logger.warning(
+ f"\n{self.__class__.__name__} has backpropagation operations that are NOT supported on CPU. If you wish "
+ "to train/fine-tune this model, you need a GPU or a TPU"
+ )
+
+ padding = int((kernel_size - 1) / 2) * dilation
+ self.padding = keras.layers.ZeroPadding2D(padding)
+
+ if out_channels % groups != 0:
+ raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups.")
+
+ self.convolution = keras.layers.Conv2D(
+ filters=out_channels,
+ kernel_size=kernel_size,
+ strides=stride,
+ padding="VALID",
+ dilation_rate=dilation,
+ groups=groups,
+ use_bias=bias,
+ name="convolution",
+ )
+
+ if use_normalization:
+ self.normalization = keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name="normalization")
+ else:
+ self.normalization = None
+
+ if use_activation:
+ if isinstance(use_activation, str):
+ self.activation = get_tf_activation(use_activation)
+ elif isinstance(config.hidden_act, str):
+ self.activation = get_tf_activation(config.hidden_act)
+ else:
+ self.activation = config.hidden_act
+ else:
+ self.activation = None
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ padded_features = self.padding(features)
+ features = self.convolution(padded_features)
+ if self.normalization is not None:
+ features = self.normalization(features, training=training)
+ if self.activation is not None:
+ features = self.activation(features)
+ return features
+
+ 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.in_channels])
+ if getattr(self, "normalization", None) is not None:
+ if hasattr(self.normalization, "name"):
+ with tf.name_scope(self.normalization.name):
+ self.normalization.build([None, None, None, self.out_channels])
+
+
+class TFMobileViTInvertedResidual(keras.layers.Layer):
+ """
+ Inverted residual block (MobileNetv2): https://arxiv.org/abs/1801.04381
+ """
+
+ def __init__(
+ self, config: MobileViTConfig, in_channels: int, out_channels: int, stride: int, dilation: int = 1, **kwargs
+ ) -> None:
+ super().__init__(**kwargs)
+ expanded_channels = make_divisible(int(round(in_channels * config.expand_ratio)), 8)
+
+ if stride not in [1, 2]:
+ raise ValueError(f"Invalid stride {stride}.")
+
+ self.use_residual = (stride == 1) and (in_channels == out_channels)
+
+ self.expand_1x1 = TFMobileViTConvLayer(
+ config, in_channels=in_channels, out_channels=expanded_channels, kernel_size=1, name="expand_1x1"
+ )
+
+ self.conv_3x3 = TFMobileViTConvLayer(
+ config,
+ in_channels=expanded_channels,
+ out_channels=expanded_channels,
+ kernel_size=3,
+ stride=stride,
+ groups=expanded_channels,
+ dilation=dilation,
+ name="conv_3x3",
+ )
+
+ self.reduce_1x1 = TFMobileViTConvLayer(
+ config,
+ in_channels=expanded_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ use_activation=False,
+ name="reduce_1x1",
+ )
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ residual = features
+
+ features = self.expand_1x1(features, training=training)
+ features = self.conv_3x3(features, training=training)
+ features = self.reduce_1x1(features, training=training)
+
+ return residual + features if self.use_residual else features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "expand_1x1", None) is not None:
+ with tf.name_scope(self.expand_1x1.name):
+ self.expand_1x1.build(None)
+ if getattr(self, "conv_3x3", None) is not None:
+ with tf.name_scope(self.conv_3x3.name):
+ self.conv_3x3.build(None)
+ if getattr(self, "reduce_1x1", None) is not None:
+ with tf.name_scope(self.reduce_1x1.name):
+ self.reduce_1x1.build(None)
+
+
+class TFMobileViTMobileNetLayer(keras.layers.Layer):
+ def __init__(
+ self,
+ config: MobileViTConfig,
+ in_channels: int,
+ out_channels: int,
+ stride: int = 1,
+ num_stages: int = 1,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+
+ self.layers = []
+ for i in range(num_stages):
+ layer = TFMobileViTInvertedResidual(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=stride if i == 0 else 1,
+ name=f"layer.{i}",
+ )
+ self.layers.append(layer)
+ in_channels = out_channels
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ for layer_module in self.layers:
+ features = layer_module(features, training=training)
+ return features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "layers", None) is not None:
+ for layer_module in self.layers:
+ with tf.name_scope(layer_module.name):
+ layer_module.build(None)
+
+
+class TFMobileViTSelfAttention(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+
+ if hidden_size % config.num_attention_heads != 0:
+ raise ValueError(
+ f"The hidden size {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(hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ scale = tf.cast(self.attention_head_size, dtype=tf.float32)
+ self.scale = tf.math.sqrt(scale)
+
+ self.query = keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="query")
+ self.key = keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="key")
+ self.value = keras.layers.Dense(self.all_head_size, use_bias=config.qkv_bias, name="value")
+
+ self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
+ self.hidden_size = hidden_size
+
+ def transpose_for_scores(self, x: tf.Tensor) -> tf.Tensor:
+ batch_size = tf.shape(x)[0]
+ x = tf.reshape(x, shape=(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: tf.Tensor, training: bool = False) -> tf.Tensor:
+ batch_size = tf.shape(hidden_states)[0]
+
+ 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(self.query(hidden_states))
+
+ # 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)
+ attention_scores = attention_scores / self.scale
+
+ # 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)
+
+ context_layer = tf.matmul(attention_probs, value_layer)
+
+ context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
+ context_layer = tf.reshape(context_layer, shape=(batch_size, -1, self.all_head_size))
+ return context_layer
+
+ 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.hidden_size])
+ if getattr(self, "key", None) is not None:
+ with tf.name_scope(self.key.name):
+ self.key.build([None, None, self.hidden_size])
+ if getattr(self, "value", None) is not None:
+ with tf.name_scope(self.value.name):
+ self.value.build([None, None, self.hidden_size])
+
+
+class TFMobileViTSelfOutput(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(hidden_size, name="dense")
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
+ self.hidden_size = hidden_size
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(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.hidden_size])
+
+
+class TFMobileViTAttention(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.attention = TFMobileViTSelfAttention(config, hidden_size, name="attention")
+ self.dense_output = TFMobileViTSelfOutput(config, hidden_size, name="output")
+
+ def prune_heads(self, heads):
+ raise NotImplementedError
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ self_outputs = self.attention(hidden_states, training=training)
+ attention_output = self.dense_output(self_outputs, 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 TFMobileViTIntermediate(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(intermediate_size, 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.hidden_size = hidden_size
+
+ def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
+ 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.hidden_size])
+
+
+class TFMobileViTOutput(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(hidden_size, name="dense")
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
+ self.intermediate_size = intermediate_size
+
+ def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(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.intermediate_size])
+
+
+class TFMobileViTTransformerLayer(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, intermediate_size: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.attention = TFMobileViTAttention(config, hidden_size, name="attention")
+ self.intermediate = TFMobileViTIntermediate(config, hidden_size, intermediate_size, name="intermediate")
+ self.mobilevit_output = TFMobileViTOutput(config, hidden_size, intermediate_size, 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.hidden_size = hidden_size
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ attention_output = self.attention(self.layernorm_before(hidden_states), training=training)
+ hidden_states = attention_output + hidden_states
+
+ layer_output = self.layernorm_after(hidden_states)
+ layer_output = self.intermediate(layer_output)
+ layer_output = self.mobilevit_output(layer_output, hidden_states, 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, "mobilevit_output", None) is not None:
+ with tf.name_scope(self.mobilevit_output.name):
+ self.mobilevit_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.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.hidden_size])
+
+
+class TFMobileViTTransformer(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, hidden_size: int, num_stages: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+
+ self.layers = []
+ for i in range(num_stages):
+ transformer_layer = TFMobileViTTransformerLayer(
+ config,
+ hidden_size=hidden_size,
+ intermediate_size=int(hidden_size * config.mlp_ratio),
+ name=f"layer.{i}",
+ )
+ self.layers.append(transformer_layer)
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ for layer_module in self.layers:
+ hidden_states = layer_module(hidden_states, training=training)
+ 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_module in self.layers:
+ with tf.name_scope(layer_module.name):
+ layer_module.build(None)
+
+
+class TFMobileViTLayer(keras.layers.Layer):
+ """
+ MobileViT block: https://arxiv.org/abs/2110.02178
+ """
+
+ def __init__(
+ self,
+ config: MobileViTConfig,
+ in_channels: int,
+ out_channels: int,
+ stride: int,
+ hidden_size: int,
+ num_stages: int,
+ dilation: int = 1,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.patch_width = config.patch_size
+ self.patch_height = config.patch_size
+
+ if stride == 2:
+ self.downsampling_layer = TFMobileViTInvertedResidual(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=stride if dilation == 1 else 1,
+ dilation=dilation // 2 if dilation > 1 else 1,
+ name="downsampling_layer",
+ )
+ in_channels = out_channels
+ else:
+ self.downsampling_layer = None
+
+ self.conv_kxk = TFMobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=config.conv_kernel_size,
+ name="conv_kxk",
+ )
+
+ self.conv_1x1 = TFMobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=hidden_size,
+ kernel_size=1,
+ use_normalization=False,
+ use_activation=False,
+ name="conv_1x1",
+ )
+
+ self.transformer = TFMobileViTTransformer(
+ config, hidden_size=hidden_size, num_stages=num_stages, name="transformer"
+ )
+
+ self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
+
+ self.conv_projection = TFMobileViTConvLayer(
+ config, in_channels=hidden_size, out_channels=in_channels, kernel_size=1, name="conv_projection"
+ )
+
+ self.fusion = TFMobileViTConvLayer(
+ config,
+ in_channels=2 * in_channels,
+ out_channels=in_channels,
+ kernel_size=config.conv_kernel_size,
+ name="fusion",
+ )
+ self.hidden_size = hidden_size
+
+ def unfolding(self, features: tf.Tensor) -> Tuple[tf.Tensor, Dict]:
+ patch_width, patch_height = self.patch_width, self.patch_height
+ patch_area = tf.cast(patch_width * patch_height, "int32")
+
+ batch_size = tf.shape(features)[0]
+ orig_height = tf.shape(features)[1]
+ orig_width = tf.shape(features)[2]
+ channels = tf.shape(features)[3]
+
+ new_height = tf.cast(tf.math.ceil(orig_height / patch_height) * patch_height, "int32")
+ new_width = tf.cast(tf.math.ceil(orig_width / patch_width) * patch_width, "int32")
+
+ interpolate = new_width != orig_width or new_height != orig_height
+ if interpolate:
+ # Note: Padding can be done, but then it needs to be handled in attention function.
+ features = tf.image.resize(features, size=(new_height, new_width), method="bilinear")
+
+ # number of patches along width and height
+ num_patch_width = new_width // patch_width
+ num_patch_height = new_height // patch_height
+ num_patches = num_patch_height * num_patch_width
+
+ # convert from shape (batch_size, orig_height, orig_width, channels)
+ # to the shape (batch_size * patch_area, num_patches, channels)
+ features = tf.transpose(features, [0, 3, 1, 2])
+ patches = tf.reshape(
+ features, (batch_size * channels * num_patch_height, patch_height, num_patch_width, patch_width)
+ )
+ patches = tf.transpose(patches, [0, 2, 1, 3])
+ patches = tf.reshape(patches, (batch_size, channels, num_patches, patch_area))
+ patches = tf.transpose(patches, [0, 3, 2, 1])
+ patches = tf.reshape(patches, (batch_size * patch_area, num_patches, channels))
+
+ info_dict = {
+ "orig_size": (orig_height, orig_width),
+ "batch_size": batch_size,
+ "channels": channels,
+ "interpolate": interpolate,
+ "num_patches": num_patches,
+ "num_patches_width": num_patch_width,
+ "num_patches_height": num_patch_height,
+ }
+ return patches, info_dict
+
+ def folding(self, patches: tf.Tensor, info_dict: Dict) -> tf.Tensor:
+ patch_width, patch_height = self.patch_width, self.patch_height
+ patch_area = int(patch_width * patch_height)
+
+ batch_size = info_dict["batch_size"]
+ channels = info_dict["channels"]
+ num_patches = info_dict["num_patches"]
+ num_patch_height = info_dict["num_patches_height"]
+ num_patch_width = info_dict["num_patches_width"]
+
+ # convert from shape (batch_size * patch_area, num_patches, channels)
+ # back to shape (batch_size, channels, orig_height, orig_width)
+ features = tf.reshape(patches, (batch_size, patch_area, num_patches, -1))
+ features = tf.transpose(features, perm=(0, 3, 2, 1))
+ features = tf.reshape(
+ features, (batch_size * channels * num_patch_height, num_patch_width, patch_height, patch_width)
+ )
+ features = tf.transpose(features, perm=(0, 2, 1, 3))
+ features = tf.reshape(
+ features, (batch_size, channels, num_patch_height * patch_height, num_patch_width * patch_width)
+ )
+ features = tf.transpose(features, perm=(0, 2, 3, 1))
+
+ if info_dict["interpolate"]:
+ features = tf.image.resize(features, size=info_dict["orig_size"], method="bilinear")
+
+ return features
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ # reduce spatial dimensions if needed
+ if self.downsampling_layer:
+ features = self.downsampling_layer(features, training=training)
+
+ residual = features
+
+ # local representation
+ features = self.conv_kxk(features, training=training)
+ features = self.conv_1x1(features, training=training)
+
+ # convert feature map to patches
+ patches, info_dict = self.unfolding(features)
+
+ # learn global representations
+ patches = self.transformer(patches, training=training)
+ patches = self.layernorm(patches)
+
+ # convert patches back to feature maps
+ features = self.folding(patches, info_dict)
+
+ features = self.conv_projection(features, training=training)
+ features = self.fusion(tf.concat([residual, features], axis=-1), training=training)
+ return features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "conv_kxk", None) is not None:
+ with tf.name_scope(self.conv_kxk.name):
+ self.conv_kxk.build(None)
+ if getattr(self, "conv_1x1", None) is not None:
+ with tf.name_scope(self.conv_1x1.name):
+ self.conv_1x1.build(None)
+ if getattr(self, "transformer", None) is not None:
+ with tf.name_scope(self.transformer.name):
+ self.transformer.build(None)
+ if getattr(self, "layernorm", None) is not None:
+ with tf.name_scope(self.layernorm.name):
+ self.layernorm.build([None, None, self.hidden_size])
+ if getattr(self, "conv_projection", None) is not None:
+ with tf.name_scope(self.conv_projection.name):
+ self.conv_projection.build(None)
+ if getattr(self, "fusion", None) is not None:
+ with tf.name_scope(self.fusion.name):
+ self.fusion.build(None)
+ if getattr(self, "downsampling_layer", None) is not None:
+ with tf.name_scope(self.downsampling_layer.name):
+ self.downsampling_layer.build(None)
+
+
+class TFMobileViTEncoder(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.config = config
+
+ self.layers = []
+
+ # segmentation architectures like DeepLab and PSPNet modify the strides
+ # of the classification backbones
+ dilate_layer_4 = dilate_layer_5 = False
+ if config.output_stride == 8:
+ dilate_layer_4 = True
+ dilate_layer_5 = True
+ elif config.output_stride == 16:
+ dilate_layer_5 = True
+
+ dilation = 1
+
+ layer_1 = TFMobileViTMobileNetLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[0],
+ out_channels=config.neck_hidden_sizes[1],
+ stride=1,
+ num_stages=1,
+ name="layer.0",
+ )
+ self.layers.append(layer_1)
+
+ layer_2 = TFMobileViTMobileNetLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[1],
+ out_channels=config.neck_hidden_sizes[2],
+ stride=2,
+ num_stages=3,
+ name="layer.1",
+ )
+ self.layers.append(layer_2)
+
+ layer_3 = TFMobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[2],
+ out_channels=config.neck_hidden_sizes[3],
+ stride=2,
+ hidden_size=config.hidden_sizes[0],
+ num_stages=2,
+ name="layer.2",
+ )
+ self.layers.append(layer_3)
+
+ if dilate_layer_4:
+ dilation *= 2
+
+ layer_4 = TFMobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[3],
+ out_channels=config.neck_hidden_sizes[4],
+ stride=2,
+ hidden_size=config.hidden_sizes[1],
+ num_stages=4,
+ dilation=dilation,
+ name="layer.3",
+ )
+ self.layers.append(layer_4)
+
+ if dilate_layer_5:
+ dilation *= 2
+
+ layer_5 = TFMobileViTLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[4],
+ out_channels=config.neck_hidden_sizes[5],
+ stride=2,
+ hidden_size=config.hidden_sizes[2],
+ num_stages=3,
+ dilation=dilation,
+ name="layer.4",
+ )
+ self.layers.append(layer_5)
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ training: bool = False,
+ ) -> Union[tuple, TFBaseModelOutput]:
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, layer_module in enumerate(self.layers):
+ hidden_states = layer_module(hidden_states, training=training)
+
+ 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):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "layers", None) is not None:
+ for layer_module in self.layers:
+ with tf.name_scope(layer_module.name):
+ layer_module.build(None)
+
+
+@keras_serializable
+class TFMobileViTMainLayer(keras.layers.Layer):
+ config_class = MobileViTConfig
+
+ def __init__(self, config: MobileViTConfig, expand_output: bool = True, **kwargs):
+ super().__init__(**kwargs)
+ self.config = config
+ self.expand_output = expand_output
+
+ self.conv_stem = TFMobileViTConvLayer(
+ config,
+ in_channels=config.num_channels,
+ out_channels=config.neck_hidden_sizes[0],
+ kernel_size=3,
+ stride=2,
+ name="conv_stem",
+ )
+
+ self.encoder = TFMobileViTEncoder(config, name="encoder")
+
+ if self.expand_output:
+ self.conv_1x1_exp = TFMobileViTConvLayer(
+ config,
+ in_channels=config.neck_hidden_sizes[5],
+ out_channels=config.neck_hidden_sizes[6],
+ kernel_size=1,
+ name="conv_1x1_exp",
+ )
+
+ self.pooler = keras.layers.GlobalAveragePooling2D(data_format="channels_first", name="pooler")
+
+ 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,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ training: bool = False,
+ ) -> Union[Tuple[tf.Tensor], TFBaseModelOutputWithPooling]:
+ 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
+
+ # 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))
+
+ embedding_output = self.conv_stem(pixel_values, training=training)
+
+ encoder_outputs = self.encoder(
+ embedding_output, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training
+ )
+
+ if self.expand_output:
+ last_hidden_state = self.conv_1x1_exp(encoder_outputs[0])
+
+ # Change to NCHW output format to have uniformity in the modules
+ last_hidden_state = tf.transpose(last_hidden_state, perm=[0, 3, 1, 2])
+
+ # global average pooling: (batch_size, channels, height, width) -> (batch_size, channels)
+ pooled_output = self.pooler(last_hidden_state)
+ else:
+ last_hidden_state = encoder_outputs[0]
+ # Change to NCHW output format to have uniformity in the modules
+ last_hidden_state = tf.transpose(last_hidden_state, perm=[0, 3, 1, 2])
+ pooled_output = None
+
+ if not return_dict:
+ output = (last_hidden_state, pooled_output) if pooled_output is not None else (last_hidden_state,)
+
+ # Change to NCHW output format to have uniformity in the modules
+ if not self.expand_output:
+ remaining_encoder_outputs = encoder_outputs[1:]
+ remaining_encoder_outputs = tuple(
+ [tf.transpose(h, perm=(0, 3, 1, 2)) for h in remaining_encoder_outputs[0]]
+ )
+ remaining_encoder_outputs = (remaining_encoder_outputs,)
+ return output + remaining_encoder_outputs
+ else:
+ return output + encoder_outputs[1:]
+
+ # 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]])
+
+ 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, "conv_stem", None) is not None:
+ with tf.name_scope(self.conv_stem.name):
+ self.conv_stem.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, None, None, None])
+ if getattr(self, "conv_1x1_exp", None) is not None:
+ with tf.name_scope(self.conv_1x1_exp.name):
+ self.conv_1x1_exp.build(None)
+
+
+class TFMobileViTPreTrainedModel(TFPreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MobileViTConfig
+ base_model_prefix = "mobilevit"
+ main_input_name = "pixel_values"
+
+
+MOBILEVIT_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 ([`MobileViTConfig`]): 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.
+"""
+
+MOBILEVIT_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
+ [`MobileViTImageProcessor.__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 MobileViT model outputting raw hidden-states without any specific head on top.",
+ MOBILEVIT_START_DOCSTRING,
+)
+class TFMobileViTModel(TFMobileViTPreTrainedModel):
+ def __init__(self, config: MobileViTConfig, expand_output: bool = True, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.config = config
+ self.expand_output = expand_output
+
+ self.mobilevit = TFMobileViTMainLayer(config, expand_output=expand_output, name="mobilevit")
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(MOBILEVIT_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,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ training: bool = False,
+ ) -> Union[Tuple[tf.Tensor], TFBaseModelOutputWithPooling]:
+ output = self.mobilevit(pixel_values, output_hidden_states, return_dict, training=training)
+ return output
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "mobilevit", None) is not None:
+ with tf.name_scope(self.mobilevit.name):
+ self.mobilevit.build(None)
+
+
+@add_start_docstrings(
+ """
+ MobileViT model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """,
+ MOBILEVIT_START_DOCSTRING,
+)
+class TFMobileViTForImageClassification(TFMobileViTPreTrainedModel, TFSequenceClassificationLoss):
+ def __init__(self, config: MobileViTConfig, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+
+ self.num_labels = config.num_labels
+ self.mobilevit = TFMobileViTMainLayer(config, name="mobilevit")
+
+ # Classifier head
+ self.dropout = keras.layers.Dropout(config.classifier_dropout_prob)
+ self.classifier = (
+ keras.layers.Dense(config.num_labels, name="classifier") if config.num_labels > 0 else tf.identity
+ )
+ self.config = config
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(MOBILEVIT_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: tf.Tensor | None = None,
+ output_hidden_states: Optional[bool] = None,
+ labels: tf.Tensor | None = None,
+ return_dict: Optional[bool] = None,
+ training: Optional[bool] = False,
+ ) -> Union[tuple, TFImageClassifierOutputWithNoAttention]:
+ 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.mobilevit(
+ 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(self.dropout(pooled_output, training=training))
+ 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, "mobilevit", None) is not None:
+ with tf.name_scope(self.mobilevit.name):
+ self.mobilevit.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.neck_hidden_sizes[-1]])
+
+
+class TFMobileViTASPPPooling(keras.layers.Layer):
+ def __init__(self, config: MobileViTConfig, in_channels: int, out_channels: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+
+ self.global_pool = keras.layers.GlobalAveragePooling2D(keepdims=True, name="global_pool")
+
+ self.conv_1x1 = TFMobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ stride=1,
+ use_normalization=True,
+ use_activation="relu",
+ name="conv_1x1",
+ )
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ spatial_size = shape_list(features)[1:-1]
+ features = self.global_pool(features)
+ features = self.conv_1x1(features, training=training)
+ features = tf.image.resize(features, size=spatial_size, method="bilinear")
+ return features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "global_pool", None) is not None:
+ with tf.name_scope(self.global_pool.name):
+ self.global_pool.build([None, None, None, None])
+ if getattr(self, "conv_1x1", None) is not None:
+ with tf.name_scope(self.conv_1x1.name):
+ self.conv_1x1.build(None)
+
+
+class TFMobileViTASPP(keras.layers.Layer):
+ """
+ ASPP module defined in DeepLab papers: https://arxiv.org/abs/1606.00915, https://arxiv.org/abs/1706.05587
+ """
+
+ def __init__(self, config: MobileViTConfig, **kwargs) -> None:
+ super().__init__(**kwargs)
+
+ in_channels = config.neck_hidden_sizes[-2]
+ out_channels = config.aspp_out_channels
+
+ if len(config.atrous_rates) != 3:
+ raise ValueError("Expected 3 values for atrous_rates")
+
+ self.convs = []
+
+ in_projection = TFMobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ use_activation="relu",
+ name="convs.0",
+ )
+ self.convs.append(in_projection)
+
+ self.convs.extend(
+ [
+ TFMobileViTConvLayer(
+ config,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ dilation=rate,
+ use_activation="relu",
+ name=f"convs.{i + 1}",
+ )
+ for i, rate in enumerate(config.atrous_rates)
+ ]
+ )
+
+ pool_layer = TFMobileViTASPPPooling(
+ config, in_channels, out_channels, name=f"convs.{len(config.atrous_rates) + 1}"
+ )
+ self.convs.append(pool_layer)
+
+ self.project = TFMobileViTConvLayer(
+ config,
+ in_channels=5 * out_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ use_activation="relu",
+ name="project",
+ )
+
+ self.dropout = keras.layers.Dropout(config.aspp_dropout_prob)
+
+ def call(self, features: tf.Tensor, training: bool = False) -> tf.Tensor:
+ # since the hidden states were transposed to have `(batch_size, channels, height, width)`
+ # layout we transpose them back to have `(batch_size, height, width, channels)` layout.
+ features = tf.transpose(features, perm=[0, 2, 3, 1])
+ pyramid = []
+ for conv in self.convs:
+ pyramid.append(conv(features, training=training))
+ pyramid = tf.concat(pyramid, axis=-1)
+
+ pooled_features = self.project(pyramid, training=training)
+ pooled_features = self.dropout(pooled_features, training=training)
+ return pooled_features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "project", None) is not None:
+ with tf.name_scope(self.project.name):
+ self.project.build(None)
+ if getattr(self, "convs", None) is not None:
+ for conv in self.convs:
+ with tf.name_scope(conv.name):
+ conv.build(None)
+
+
+class TFMobileViTDeepLabV3(keras.layers.Layer):
+ """
+ DeepLabv3 architecture: https://arxiv.org/abs/1706.05587
+ """
+
+ def __init__(self, config: MobileViTConfig, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.aspp = TFMobileViTASPP(config, name="aspp")
+
+ self.dropout = keras.layers.Dropout(config.classifier_dropout_prob)
+
+ self.classifier = TFMobileViTConvLayer(
+ config,
+ in_channels=config.aspp_out_channels,
+ out_channels=config.num_labels,
+ kernel_size=1,
+ use_normalization=False,
+ use_activation=False,
+ bias=True,
+ name="classifier",
+ )
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ features = self.aspp(hidden_states[-1], training=training)
+ features = self.dropout(features, training=training)
+ features = self.classifier(features, training=training)
+ return features
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "aspp", None) is not None:
+ with tf.name_scope(self.aspp.name):
+ self.aspp.build(None)
+ if getattr(self, "classifier", None) is not None:
+ with tf.name_scope(self.classifier.name):
+ self.classifier.build(None)
+
+
+@add_start_docstrings(
+ """
+ MobileViT model with a semantic segmentation head on top, e.g. for Pascal VOC.
+ """,
+ MOBILEVIT_START_DOCSTRING,
+)
+class TFMobileViTForSemanticSegmentation(TFMobileViTPreTrainedModel):
+ def __init__(self, config: MobileViTConfig, **kwargs) -> None:
+ super().__init__(config, **kwargs)
+
+ self.num_labels = config.num_labels
+ self.mobilevit = TFMobileViTMainLayer(config, expand_output=False, name="mobilevit")
+ self.segmentation_head = TFMobileViTDeepLabV3(config, name="segmentation_head")
+
+ def hf_compute_loss(self, logits, labels):
+ # upsample logits to the images' original size
+ # `labels` is of shape (batch_size, height, width)
+ label_interp_shape = shape_list(labels)[1:]
+
+ upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
+ # compute weighted loss
+ loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")
+
+ def masked_loss(real, pred):
+ unmasked_loss = loss_fct(real, pred)
+ mask = tf.cast(real != self.config.semantic_loss_ignore_index, dtype=unmasked_loss.dtype)
+ masked_loss = unmasked_loss * mask
+ # Reduction strategy in the similar spirit with
+ # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L210
+ reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(mask)
+ return tf.reshape(reduced_masked_loss, (1,))
+
+ return masked_loss(labels, upsampled_logits)
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(MOBILEVIT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFSemanticSegmenterOutputWithNoAttention, 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: bool = False,
+ ) -> Union[tuple, TFSemanticSegmenterOutputWithNoAttention]:
+ 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, TFMobileViTForSemanticSegmentation
+ >>> 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("apple/deeplabv3-mobilevit-small")
+ >>> model = TFMobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-small")
+
+ >>> inputs = image_processor(images=image, return_tensors="tf")
+
+ >>> outputs = model(**inputs)
+
+ >>> # logits are of shape (batch_size, num_labels, height, width)
+ >>> logits = outputs.logits
+ ```"""
+ 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 labels is not None and not self.config.num_labels > 1:
+ raise ValueError("The number of labels should be greater than one")
+
+ outputs = self.mobilevit(
+ pixel_values,
+ output_hidden_states=True, # we need the intermediate hidden states
+ return_dict=return_dict,
+ training=training,
+ )
+
+ encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1]
+
+ logits = self.segmentation_head(encoder_hidden_states, training=training)
+
+ loss = None
+ if labels is not None:
+ loss = self.hf_compute_loss(logits=logits, labels=labels)
+
+ # make logits of shape (batch_size, num_labels, height, width) to
+ # keep them consistent across APIs
+ logits = tf.transpose(logits, perm=[0, 3, 1, 2])
+
+ 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 TFSemanticSegmenterOutputWithNoAttention(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states if output_hidden_states else None,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "mobilevit", None) is not None:
+ with tf.name_scope(self.mobilevit.name):
+ self.mobilevit.build(None)
+ if getattr(self, "segmentation_head", None) is not None:
+ with tf.name_scope(self.segmentation_head.name):
+ self.segmentation_head.build(None)
+
+
+__all__ = [
+ "TFMobileViTForImageClassification",
+ "TFMobileViTForSemanticSegmentation",
+ "TFMobileViTModel",
+ "TFMobileViTPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/mra/__init__.py b/lib/python3.12/site-packages/transformers/models/mra/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2963ad0c97ba8c0c5bdc32ff9adca79ffc56cde3
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mra/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_mra import *
+ from .modeling_mra import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/mra/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c9195e9f4b6ac32d1e926a2b3def43013d09d25d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mra/__pycache__/configuration_mra.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/configuration_mra.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fe7aa2b4ee65fdd917d8dd911ddb461cfc67b7e6
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/configuration_mra.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mra/__pycache__/modeling_mra.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/modeling_mra.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c5825ae0981af1e443149fa50c47f4a4b065bb97
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/mra/__pycache__/modeling_mra.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/mra/configuration_mra.py b/lib/python3.12/site-packages/transformers/models/mra/configuration_mra.py
new file mode 100644
index 0000000000000000000000000000000000000000..16b064c98f7e6a9ed4fe72d7149fb867380c904d
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mra/configuration_mra.py
@@ -0,0 +1,137 @@
+# 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.
+"""MRA model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MraConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MraModel`]. It is used to instantiate an MRA
+ 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 Mra
+ [uw-madison/mra-base-512-4](https://huggingface.co/uw-madison/mra-base-512-4) 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 50265):
+ Vocabulary size of the Mra model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`MraModel`].
+ 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.
+ 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):
+ 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.
+ 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 1):
+ The vocabulary size of the `token_type_ids` passed when calling [`MraModel`].
+ 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.
+ position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
+ Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`.
+ block_per_row (`int`, *optional*, defaults to 4):
+ Used to set the budget for the high resolution scale.
+ approx_mode (`str`, *optional*, defaults to `"full"`):
+ Controls whether both low and high resolution approximations are used. Set to `"full"` for both low and
+ high resolution and `"sparse"` for only low resolution.
+ initial_prior_first_n_blocks (`int`, *optional*, defaults to 0):
+ The initial number of blocks for which high resolution is used.
+ initial_prior_diagonal_n_blocks (`int`, *optional*, defaults to 0):
+ The number of diagonal blocks for which high resolution is used.
+
+ Example:
+
+ ```python
+ >>> from transformers import MraConfig, MraModel
+
+ >>> # Initializing a Mra uw-madison/mra-base-512-4 style configuration
+ >>> configuration = MraConfig()
+
+ >>> # Initializing a model (with random weights) from the uw-madison/mra-base-512-4 style configuration
+ >>> model = MraModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "mra"
+
+ 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=1,
+ initializer_range=0.02,
+ layer_norm_eps=1e-5,
+ position_embedding_type="absolute",
+ block_per_row=4,
+ approx_mode="full",
+ initial_prior_first_n_blocks=0,
+ initial_prior_diagonal_n_blocks=0,
+ pad_token_id=1,
+ bos_token_id=0,
+ 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.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.type_vocab_size = type_vocab_size
+ self.layer_norm_eps = layer_norm_eps
+ self.position_embedding_type = position_embedding_type
+ self.block_per_row = block_per_row
+ self.approx_mode = approx_mode
+ self.initial_prior_first_n_blocks = initial_prior_first_n_blocks
+ self.initial_prior_diagonal_n_blocks = initial_prior_diagonal_n_blocks
+
+
+__all__ = ["MraConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/mra/modeling_mra.py b/lib/python3.12/site-packages/transformers/models/mra/modeling_mra.py
new file mode 100644
index 0000000000000000000000000000000000000000..6071d6a9e9c04d3303d43be0dfffca8158de06ae
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/mra/modeling_mra.py
@@ -0,0 +1,1492 @@
+# coding=utf-8
+# Copyright 2023 University of Wisconsin-Madison 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 MRA model."""
+
+import math
+from pathlib import Path
+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.utils.cpp_extension import load
+
+from ...activations import ACT2FN
+from ...modeling_outputs import (
+ BaseModelOutputWithCrossAttentions,
+ 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_ninja_available,
+ is_torch_cuda_available,
+ logging,
+)
+from .configuration_mra import MraConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "uw-madison/mra-base-512-4"
+_CONFIG_FOR_DOC = "MraConfig"
+_TOKENIZER_FOR_DOC = "AutoTokenizer"
+
+
+mra_cuda_kernel = None
+
+
+def load_cuda_kernels():
+ global mra_cuda_kernel
+ src_folder = Path(__file__).resolve().parent.parent.parent / "kernels" / "mra"
+
+ def append_root(files):
+ return [src_folder / file for file in files]
+
+ src_files = append_root(["cuda_kernel.cu", "cuda_launch.cu", "torch_extension.cpp"])
+
+ mra_cuda_kernel = load("cuda_kernel", src_files, verbose=True)
+
+
+def sparse_max(sparse_qk_prod, indices, query_num_block, key_num_block):
+ """
+ Computes maximum values for softmax stability.
+ """
+ if len(sparse_qk_prod.size()) != 4:
+ raise ValueError("sparse_qk_prod must be a 4-dimensional tensor.")
+
+ if len(indices.size()) != 2:
+ raise ValueError("indices must be a 2-dimensional tensor.")
+
+ if sparse_qk_prod.size(2) != 32:
+ raise ValueError("The size of the second dimension of sparse_qk_prod must be 32.")
+
+ if sparse_qk_prod.size(3) != 32:
+ raise ValueError("The size of the third dimension of sparse_qk_prod must be 32.")
+
+ index_vals = sparse_qk_prod.max(dim=-2).values.transpose(-1, -2)
+ index_vals = index_vals.contiguous()
+
+ indices = indices.int()
+ indices = indices.contiguous()
+
+ max_vals, max_vals_scatter = mra_cuda_kernel.index_max(index_vals, indices, query_num_block, key_num_block)
+ max_vals_scatter = max_vals_scatter.transpose(-1, -2)[:, :, None, :]
+
+ return max_vals, max_vals_scatter
+
+
+def sparse_mask(mask, indices, block_size=32):
+ """
+ Converts attention mask to a sparse mask for high resolution logits.
+ """
+ if len(mask.size()) != 2:
+ raise ValueError("mask must be a 2-dimensional tensor.")
+
+ if len(indices.size()) != 2:
+ raise ValueError("indices must be a 2-dimensional tensor.")
+
+ if mask.shape[0] != indices.shape[0]:
+ raise ValueError("mask and indices must have the same size in the zero-th dimension.")
+
+ batch_size, seq_len = mask.shape
+ num_block = seq_len // block_size
+
+ batch_idx = torch.arange(indices.size(0), dtype=torch.long, device=indices.device)
+ mask = mask.reshape(batch_size, num_block, block_size)
+ mask = mask[batch_idx[:, None], (indices % num_block).long(), :]
+
+ return mask
+
+
+def mm_to_sparse(dense_query, dense_key, indices, block_size=32):
+ """
+ Performs Sampled Dense Matrix Multiplication.
+ """
+ batch_size, query_size, dim = dense_query.size()
+ _, key_size, dim = dense_key.size()
+
+ if query_size % block_size != 0:
+ raise ValueError("query_size (size of first dimension of dense_query) must be divisible by block_size.")
+
+ if key_size % block_size != 0:
+ raise ValueError("key_size (size of first dimension of dense_key) must be divisible by block_size.")
+
+ dense_query = dense_query.reshape(batch_size, query_size // block_size, block_size, dim).transpose(-1, -2)
+ dense_key = dense_key.reshape(batch_size, key_size // block_size, block_size, dim).transpose(-1, -2)
+
+ if len(dense_query.size()) != 4:
+ raise ValueError("dense_query must be a 4-dimensional tensor.")
+
+ if len(dense_key.size()) != 4:
+ raise ValueError("dense_key must be a 4-dimensional tensor.")
+
+ if len(indices.size()) != 2:
+ raise ValueError("indices must be a 2-dimensional tensor.")
+
+ if dense_query.size(3) != 32:
+ raise ValueError("The third dimension of dense_query must be 32.")
+
+ if dense_key.size(3) != 32:
+ raise ValueError("The third dimension of dense_key must be 32.")
+
+ dense_query = dense_query.contiguous()
+ dense_key = dense_key.contiguous()
+
+ indices = indices.int()
+ indices = indices.contiguous()
+
+ return mra_cuda_kernel.mm_to_sparse(dense_query, dense_key, indices.int())
+
+
+def sparse_dense_mm(sparse_query, indices, dense_key, query_num_block, block_size=32):
+ """
+ Performs matrix multiplication of a sparse matrix with a dense matrix.
+ """
+ batch_size, key_size, dim = dense_key.size()
+
+ if key_size % block_size != 0:
+ raise ValueError("key_size (size of first dimension of dense_key) must be divisible by block_size.")
+
+ if sparse_query.size(2) != block_size:
+ raise ValueError("The size of the second dimension of sparse_query must be equal to the block_size.")
+
+ if sparse_query.size(3) != block_size:
+ raise ValueError("The size of the third dimension of sparse_query must be equal to the block_size.")
+
+ dense_key = dense_key.reshape(batch_size, key_size // block_size, block_size, dim).transpose(-1, -2)
+
+ if len(sparse_query.size()) != 4:
+ raise ValueError("sparse_query must be a 4-dimensional tensor.")
+
+ if len(dense_key.size()) != 4:
+ raise ValueError("dense_key must be a 4-dimensional tensor.")
+
+ if len(indices.size()) != 2:
+ raise ValueError("indices must be a 2-dimensional tensor.")
+
+ if dense_key.size(3) != 32:
+ raise ValueError("The size of the third dimension of dense_key must be 32.")
+
+ sparse_query = sparse_query.contiguous()
+
+ indices = indices.int()
+ indices = indices.contiguous()
+ dense_key = dense_key.contiguous()
+
+ dense_qk_prod = mra_cuda_kernel.sparse_dense_mm(sparse_query, indices, dense_key, query_num_block)
+ dense_qk_prod = dense_qk_prod.transpose(-1, -2).reshape(batch_size, query_num_block * block_size, dim)
+ return dense_qk_prod
+
+
+def transpose_indices(indices, dim_1_block, dim_2_block):
+ return ((indices % dim_2_block) * dim_1_block + torch.div(indices, dim_2_block, rounding_mode="floor")).long()
+
+
+class MraSampledDenseMatMul(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, dense_query, dense_key, indices, block_size):
+ sparse_qk_prod = mm_to_sparse(dense_query, dense_key, indices, block_size)
+ ctx.save_for_backward(dense_query, dense_key, indices)
+ ctx.block_size = block_size
+ return sparse_qk_prod
+
+ @staticmethod
+ def backward(ctx, grad):
+ dense_query, dense_key, indices = ctx.saved_tensors
+ block_size = ctx.block_size
+ query_num_block = dense_query.size(1) // block_size
+ key_num_block = dense_key.size(1) // block_size
+ indices_T = transpose_indices(indices, query_num_block, key_num_block)
+ grad_key = sparse_dense_mm(grad.transpose(-1, -2), indices_T, dense_query, key_num_block)
+ grad_query = sparse_dense_mm(grad, indices, dense_key, query_num_block)
+ return grad_query, grad_key, None, None
+
+ @staticmethod
+ def operator_call(dense_query, dense_key, indices, block_size=32):
+ return MraSampledDenseMatMul.apply(dense_query, dense_key, indices, block_size)
+
+
+class MraSparseDenseMatMul(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, sparse_query, indices, dense_key, query_num_block):
+ sparse_qk_prod = sparse_dense_mm(sparse_query, indices, dense_key, query_num_block)
+ ctx.save_for_backward(sparse_query, indices, dense_key)
+ ctx.query_num_block = query_num_block
+ return sparse_qk_prod
+
+ @staticmethod
+ def backward(ctx, grad):
+ sparse_query, indices, dense_key = ctx.saved_tensors
+ query_num_block = ctx.query_num_block
+ key_num_block = dense_key.size(1) // sparse_query.size(-1)
+ indices_T = transpose_indices(indices, query_num_block, key_num_block)
+ grad_key = sparse_dense_mm(sparse_query.transpose(-1, -2), indices_T, grad, key_num_block)
+ grad_query = mm_to_sparse(grad, dense_key, indices)
+ return grad_query, None, grad_key, None
+
+ @staticmethod
+ def operator_call(sparse_query, indices, dense_key, query_num_block):
+ return MraSparseDenseMatMul.apply(sparse_query, indices, dense_key, query_num_block)
+
+
+class MraReduceSum:
+ @staticmethod
+ def operator_call(sparse_query, indices, query_num_block, key_num_block):
+ batch_size, num_block, block_size, _ = sparse_query.size()
+
+ if len(sparse_query.size()) != 4:
+ raise ValueError("sparse_query must be a 4-dimensional tensor.")
+
+ if len(indices.size()) != 2:
+ raise ValueError("indices must be a 2-dimensional tensor.")
+
+ _, _, block_size, _ = sparse_query.size()
+ batch_size, num_block = indices.size()
+
+ sparse_query = sparse_query.sum(dim=2).reshape(batch_size * num_block, block_size)
+
+ batch_idx = torch.arange(indices.size(0), dtype=torch.long, device=indices.device)
+ global_idxes = (
+ torch.div(indices, key_num_block, rounding_mode="floor").long() + batch_idx[:, None] * query_num_block
+ ).reshape(batch_size * num_block)
+ temp = torch.zeros(
+ (batch_size * query_num_block, block_size), dtype=sparse_query.dtype, device=sparse_query.device
+ )
+ output = temp.index_add(0, global_idxes, sparse_query).reshape(batch_size, query_num_block, block_size)
+
+ output = output.reshape(batch_size, query_num_block * block_size)
+ return output
+
+
+def get_low_resolution_logit(query, key, block_size, mask=None, value=None):
+ """
+ Compute low resolution approximation.
+ """
+ batch_size, seq_len, head_dim = query.size()
+
+ num_block_per_row = seq_len // block_size
+
+ value_hat = None
+ if mask is not None:
+ token_count = mask.reshape(batch_size, num_block_per_row, block_size).sum(dim=-1)
+ query_hat = query.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / (
+ token_count[:, :, None] + 1e-6
+ )
+ key_hat = key.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / (
+ token_count[:, :, None] + 1e-6
+ )
+ if value is not None:
+ value_hat = value.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / (
+ token_count[:, :, None] + 1e-6
+ )
+ else:
+ token_count = block_size * torch.ones(batch_size, num_block_per_row, dtype=torch.float, device=query.device)
+ query_hat = query.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2)
+ key_hat = key.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2)
+ if value is not None:
+ value_hat = value.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2)
+
+ low_resolution_logit = torch.matmul(query_hat, key_hat.transpose(-1, -2)) / math.sqrt(head_dim)
+
+ low_resolution_logit_row_max = low_resolution_logit.max(dim=-1, keepdims=True).values
+
+ if mask is not None:
+ low_resolution_logit = (
+ low_resolution_logit - 1e4 * ((token_count[:, None, :] * token_count[:, :, None]) < 0.5).float()
+ )
+
+ return low_resolution_logit, token_count, low_resolution_logit_row_max, value_hat
+
+
+def get_block_idxes(
+ low_resolution_logit, num_blocks, approx_mode, initial_prior_first_n_blocks, initial_prior_diagonal_n_blocks
+):
+ """
+ Compute the indices of the subset of components to be used in the approximation.
+ """
+ batch_size, total_blocks_per_row, _ = low_resolution_logit.shape
+
+ if initial_prior_diagonal_n_blocks > 0:
+ offset = initial_prior_diagonal_n_blocks // 2
+ temp_mask = torch.ones(total_blocks_per_row, total_blocks_per_row, device=low_resolution_logit.device)
+ diagonal_mask = torch.tril(torch.triu(temp_mask, diagonal=-offset), diagonal=offset)
+ low_resolution_logit = low_resolution_logit + diagonal_mask[None, :, :] * 5e3
+
+ if initial_prior_first_n_blocks > 0:
+ low_resolution_logit[:, :initial_prior_first_n_blocks, :] = (
+ low_resolution_logit[:, :initial_prior_first_n_blocks, :] + 5e3
+ )
+ low_resolution_logit[:, :, :initial_prior_first_n_blocks] = (
+ low_resolution_logit[:, :, :initial_prior_first_n_blocks] + 5e3
+ )
+
+ top_k_vals = torch.topk(
+ low_resolution_logit.reshape(batch_size, -1), num_blocks, dim=-1, largest=True, sorted=False
+ )
+ indices = top_k_vals.indices
+
+ if approx_mode == "full":
+ threshold = top_k_vals.values.min(dim=-1).values
+ high_resolution_mask = (low_resolution_logit >= threshold[:, None, None]).float()
+ elif approx_mode == "sparse":
+ high_resolution_mask = None
+ else:
+ raise ValueError(f"{approx_mode} is not a valid approx_model value.")
+
+ return indices, high_resolution_mask
+
+
+def mra2_attention(
+ query,
+ key,
+ value,
+ mask,
+ num_blocks,
+ approx_mode,
+ block_size=32,
+ initial_prior_first_n_blocks=0,
+ initial_prior_diagonal_n_blocks=0,
+):
+ """
+ Use Mra to approximate self-attention.
+ """
+ if mra_cuda_kernel is None:
+ return torch.zeros_like(query).requires_grad_()
+
+ batch_size, num_head, seq_len, head_dim = query.size()
+ meta_batch = batch_size * num_head
+
+ if seq_len % block_size != 0:
+ raise ValueError("sequence length must be divisible by the block_size.")
+
+ num_block_per_row = seq_len // block_size
+
+ query = query.reshape(meta_batch, seq_len, head_dim)
+ key = key.reshape(meta_batch, seq_len, head_dim)
+ value = value.reshape(meta_batch, seq_len, head_dim)
+
+ if mask is not None:
+ query = query * mask[:, :, None]
+ key = key * mask[:, :, None]
+ value = value * mask[:, :, None]
+
+ if approx_mode == "full":
+ low_resolution_logit, token_count, low_resolution_logit_row_max, value_hat = get_low_resolution_logit(
+ query, key, block_size, mask, value
+ )
+ elif approx_mode == "sparse":
+ with torch.no_grad():
+ low_resolution_logit, token_count, low_resolution_logit_row_max, _ = get_low_resolution_logit(
+ query, key, block_size, mask
+ )
+ else:
+ raise Exception('approx_mode must be "full" or "sparse"')
+
+ with torch.no_grad():
+ low_resolution_logit_normalized = low_resolution_logit - low_resolution_logit_row_max
+ indices, high_resolution_mask = get_block_idxes(
+ low_resolution_logit_normalized,
+ num_blocks,
+ approx_mode,
+ initial_prior_first_n_blocks,
+ initial_prior_diagonal_n_blocks,
+ )
+
+ high_resolution_logit = MraSampledDenseMatMul.operator_call(
+ query, key, indices, block_size=block_size
+ ) / math.sqrt(head_dim)
+ max_vals, max_vals_scatter = sparse_max(high_resolution_logit, indices, num_block_per_row, num_block_per_row)
+ high_resolution_logit = high_resolution_logit - max_vals_scatter
+ if mask is not None:
+ high_resolution_logit = high_resolution_logit - 1e4 * (1 - sparse_mask(mask, indices)[:, :, :, None])
+ high_resolution_attn = torch.exp(high_resolution_logit)
+ high_resolution_attn_out = MraSparseDenseMatMul.operator_call(
+ high_resolution_attn, indices, value, num_block_per_row
+ )
+ high_resolution_normalizer = MraReduceSum.operator_call(
+ high_resolution_attn, indices, num_block_per_row, num_block_per_row
+ )
+
+ if approx_mode == "full":
+ low_resolution_attn = (
+ torch.exp(low_resolution_logit - low_resolution_logit_row_max - 1e4 * high_resolution_mask)
+ * token_count[:, None, :]
+ )
+
+ low_resolution_attn_out = (
+ torch.matmul(low_resolution_attn, value_hat)[:, :, None, :]
+ .repeat(1, 1, block_size, 1)
+ .reshape(meta_batch, seq_len, head_dim)
+ )
+ low_resolution_normalizer = (
+ low_resolution_attn.sum(dim=-1)[:, :, None].repeat(1, 1, block_size).reshape(meta_batch, seq_len)
+ )
+
+ log_correction = low_resolution_logit_row_max.repeat(1, 1, block_size).reshape(meta_batch, seq_len) - max_vals
+ if mask is not None:
+ log_correction = log_correction * mask
+
+ low_resolution_corr = torch.exp(log_correction * (log_correction <= 0).float())
+ low_resolution_attn_out = low_resolution_attn_out * low_resolution_corr[:, :, None]
+ low_resolution_normalizer = low_resolution_normalizer * low_resolution_corr
+
+ high_resolution_corr = torch.exp(-log_correction * (log_correction > 0).float())
+ high_resolution_attn_out = high_resolution_attn_out * high_resolution_corr[:, :, None]
+ high_resolution_normalizer = high_resolution_normalizer * high_resolution_corr
+
+ context_layer = (high_resolution_attn_out + low_resolution_attn_out) / (
+ high_resolution_normalizer[:, :, None] + low_resolution_normalizer[:, :, None] + 1e-6
+ )
+
+ elif approx_mode == "sparse":
+ context_layer = high_resolution_attn_out / (high_resolution_normalizer[:, :, None] + 1e-6)
+ else:
+ raise Exception('config.approx_mode must be "full" or "sparse"')
+
+ if mask is not None:
+ context_layer = context_layer * mask[:, :, None]
+
+ context_layer = context_layer.reshape(batch_size, num_head, seq_len, head_dim)
+
+ return context_layer
+
+
+class MraEmbeddings(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 + 2, 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.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)) + 2)
+ 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, device=self.position_ids.device),
+ 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
+ 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 MraSelfAttention(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})"
+ )
+
+ kernel_loaded = mra_cuda_kernel 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}")
+
+ 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 if position_embedding_type is not None else config.position_embedding_type
+ )
+
+ self.num_block = (config.max_position_embeddings // 32) * config.block_per_row
+ self.num_block = min(self.num_block, int((config.max_position_embeddings // 32) ** 2))
+
+ self.approx_mode = config.approx_mode
+ self.initial_prior_first_n_blocks = config.initial_prior_first_n_blocks
+ self.initial_prior_diagonal_n_blocks = config.initial_prior_diagonal_n_blocks
+
+ def transpose_for_scores(self, layer):
+ new_layer_shape = layer.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+ layer = layer.view(*new_layer_shape)
+ return layer.permute(0, 2, 1, 3)
+
+ def forward(self, hidden_states, attention_mask=None):
+ 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)
+
+ batch_size, num_heads, seq_len, head_dim = query_layer.size()
+
+ # revert changes made by get_extended_attention_mask
+ attention_mask = 1.0 + attention_mask / 10000.0
+ attention_mask = (
+ attention_mask.squeeze().repeat(1, num_heads, 1).reshape(batch_size * num_heads, seq_len).int()
+ )
+
+ # The CUDA kernels are most efficient with inputs whose size is a multiple of a GPU's warp size (32). Inputs
+ # smaller than this are padded with zeros.
+ gpu_warp_size = 32
+
+ if head_dim < gpu_warp_size:
+ pad_size = batch_size, num_heads, seq_len, gpu_warp_size - head_dim
+
+ query_layer = torch.cat([query_layer, torch.zeros(pad_size, device=query_layer.device)], dim=-1)
+ key_layer = torch.cat([key_layer, torch.zeros(pad_size, device=key_layer.device)], dim=-1)
+ value_layer = torch.cat([value_layer, torch.zeros(pad_size, device=value_layer.device)], dim=-1)
+
+ context_layer = mra2_attention(
+ query_layer.float(),
+ key_layer.float(),
+ value_layer.float(),
+ attention_mask.float(),
+ self.num_block,
+ approx_mode=self.approx_mode,
+ initial_prior_first_n_blocks=self.initial_prior_first_n_blocks,
+ initial_prior_diagonal_n_blocks=self.initial_prior_diagonal_n_blocks,
+ )
+
+ if head_dim < gpu_warp_size:
+ context_layer = context_layer[:, :, :, :head_dim]
+
+ context_layer = context_layer.reshape(batch_size, num_heads, seq_len, head_dim)
+
+ 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,)
+
+ return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class MraSelfOutput(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 MraAttention(nn.Module):
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ self.self = MraSelfAttention(config, position_embedding_type=position_embedding_type)
+ self.output = MraSelfOutput(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):
+ self_outputs = self.self(hidden_states, attention_mask)
+ 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 MraIntermediate(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 MraOutput(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 MraLayer(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 = MraAttention(config)
+ self.add_cross_attention = config.add_cross_attention
+ self.intermediate = MraIntermediate(config)
+ self.output = MraOutput(config)
+
+ def forward(self, hidden_states, attention_mask=None):
+ self_attention_outputs = self.attention(hidden_states, attention_mask)
+ 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 MraEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([MraLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ 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,
+ attention_mask,
+ )
+ else:
+ layer_outputs = layer_module(hidden_states, attention_mask)
+
+ 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 BaseModelOutputWithCrossAttentions(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform
+class MraPredictionHeadTransform(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->Mra
+class MraLMPredictionHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.transform = MraPredictionHeadTransform(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 _tie_weights(self):
+ 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->Mra
+class MraOnlyMLMHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = MraLMPredictionHead(config)
+
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+ prediction_scores = self.predictions(sequence_output)
+ return prediction_scores
+
+
+# Copied from transformers.models.yoso.modeling_yoso.YosoPreTrainedModel with Yoso->Mra,yoso->mra
+class MraPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MraConfig
+ base_model_prefix = "mra"
+ 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)
+
+
+MRA_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 ([`MraConfig`]): 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.
+"""
+
+MRA_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_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 MRA Model transformer outputting raw hidden-states without any specific head on top.",
+ MRA_START_DOCSTRING,
+)
+class MraModel(MraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = MraEmbeddings(config)
+ self.encoder = MraEncoder(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(MRA_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,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithCrossAttentions]:
+ 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(((batch_size, seq_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)
+
+ # 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,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ 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 BaseModelOutputWithCrossAttentions(
+ last_hidden_state=sequence_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_attentions,
+ )
+
+
+@add_start_docstrings("""MRA Model with a `language modeling` head on top.""", MRA_START_DOCSTRING)
+class MraForMaskedLM(MraPreTrainedModel):
+ _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.mra = MraModel(config)
+ self.cls = MraOnlyMLMHead(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
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @add_start_docstrings_to_model_forward(MRA_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,
+ 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_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.mra(
+ 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_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.yoso.modeling_yoso.YosoClassificationHead with Yoso->Mra
+class MraClassificationHead(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)
+
+ self.config = config
+
+ def forward(self, features, **kwargs):
+ x = features[:, 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(
+ """MRA Model transformer with a sequence classification/regression head on top (a linear layer on top of
+ the pooled output) e.g. for GLUE tasks.""",
+ MRA_START_DOCSTRING,
+)
+class MraForSequenceClassification(MraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.mra = MraModel(config)
+ self.classifier = MraClassificationHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(MRA_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,
+ head_mask: 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.mra(
+ 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_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(
+ """MRA 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.""",
+ MRA_START_DOCSTRING,
+)
+class MraForMultipleChoice(MraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.mra = MraModel(config)
+ self.pre_classifier = nn.Linear(config.hidden_size, config.hidden_size)
+ self.classifier = nn.Linear(config.hidden_size, 1)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(MRA_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_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.mra(
+ 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_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)
+ 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(
+ """MRA 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.""",
+ MRA_START_DOCSTRING,
+)
+class MraForTokenClassification(MraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.mra = MraModel(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(MRA_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,
+ head_mask: 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.mra(
+ 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_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
+ if attention_mask is not None:
+ active_loss = attention_mask.view(-1) == 1
+ active_logits = logits.view(-1, self.num_labels)
+ active_labels = torch.where(
+ active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+ )
+ loss = loss_fct(active_logits, active_labels)
+ else:
+ 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(
+ """MRA 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`).""",
+ MRA_START_DOCSTRING,
+)
+class MraForQuestionAnswering(MraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ config.num_labels = 2
+ self.num_labels = config.num_labels
+
+ self.mra = MraModel(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(MRA_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,
+ 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_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.mra(
+ 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_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)
+ 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 = 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,
+ )
+
+
+__all__ = [
+ "MraForMaskedLM",
+ "MraForMultipleChoice",
+ "MraForQuestionAnswering",
+ "MraForSequenceClassification",
+ "MraForTokenClassification",
+ "MraLayer",
+ "MraModel",
+ "MraPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py b/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..880274309cbab486a90df05548a0e4d3f2ea0925
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py
@@ -0,0 +1,28 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_musicgen import *
+ from .modeling_musicgen import *
+ from .processing_musicgen import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..85361ee33ff9f856a10d6b5c32b73a5738958a1e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/configuration_musicgen.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/configuration_musicgen.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..db2224e9f13cf00f73517500edac1f47e09ebc9b
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/configuration_musicgen.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/processing_musicgen.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/processing_musicgen.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b5172e61a8e542cc8e1196b202de2bf2a66bf621
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen/__pycache__/processing_musicgen.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py b/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c38caf20dc4136e794e5762bf0b58847817dd19
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py
@@ -0,0 +1,247 @@
+# 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.
+"""MusicGen model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ..auto.configuration_auto import AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class MusicgenDecoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of an [`MusicgenDecoder`]. It is used to instantiate a
+ MusicGen decoder according to the specified arguments, defining the model architecture. Instantiating a
+ configuration with the defaults will yield a similar configuration to that of the MusicGen
+ [facebook/musicgen-small](https://huggingface.co/facebook/musicgen-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 2048):
+ Vocabulary size of the MusicgenDecoder model. Defines the number of different tokens that can be
+ represented by the `inputs_ids` passed when calling [`MusicgenDecoder`].
+ hidden_size (`int`, *optional*, defaults to 1024):
+ Dimensionality of the layers and the pooler layer.
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of decoder layers.
+ num_attention_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads for each attention layer in the Transformer block.
+ ffn_dim (`int`, *optional*, defaults to 4096):
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer block.
+ activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the decoder 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, text_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 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_factor (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ 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.
+ scale_embedding (`bool`, *optional*, defaults to `False`):
+ Scale embeddings by diving by sqrt(hidden_size).
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether the model should return the last key/values attentions (not used by all models)
+ num_codebooks (`int`, *optional*, defaults to 4):
+ The number of parallel codebooks forwarded to the model.
+ tie_word_embeddings(`bool`, *optional*, defaults to `False`):
+ Whether input and output word embeddings should be tied.
+ audio_channels (`int`, *optional*, defaults to 1
+ Number of channels in the audio data. Either 1 for mono or 2 for stereo. Stereo models generate a separate
+ audio stream for the left/right output channels. Mono models generate a single audio stream output.
+ """
+
+ model_type = "musicgen_decoder"
+ base_config_key = "decoder_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=2048,
+ max_position_embeddings=2048,
+ num_hidden_layers=24,
+ ffn_dim=4096,
+ num_attention_heads=16,
+ layerdrop=0.0,
+ use_cache=True,
+ activation_function="gelu",
+ hidden_size=1024,
+ dropout=0.1,
+ attention_dropout=0.0,
+ activation_dropout=0.0,
+ initializer_factor=0.02,
+ scale_embedding=False,
+ num_codebooks=4,
+ audio_channels=1,
+ pad_token_id=2048,
+ bos_token_id=2048,
+ eos_token_id=None,
+ tie_word_embeddings=False,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.ffn_dim = ffn_dim
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.dropout = dropout
+ self.attention_dropout = attention_dropout
+ self.activation_dropout = activation_dropout
+ self.activation_function = activation_function
+ self.initializer_factor = initializer_factor
+ self.layerdrop = layerdrop
+ self.use_cache = use_cache
+ self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
+ self.num_codebooks = num_codebooks
+
+ if audio_channels not in [1, 2]:
+ raise ValueError(f"Expected 1 (mono) or 2 (stereo) audio channels, got {audio_channels} channels.")
+ self.audio_channels = audio_channels
+
+ 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,
+ )
+
+
+class MusicgenConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MusicgenModel`]. It is used to instantiate a
+ MusicGen model according to the specified arguments, defining the text encoder, audio encoder and MusicGen 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:
+
+ - **text_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+ defines the text encoder config.
+ - **audio_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+ defines the audio encoder config.
+ - **decoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that defines
+ the decoder config.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... MusicgenConfig,
+ ... MusicgenDecoderConfig,
+ ... T5Config,
+ ... EncodecConfig,
+ ... MusicgenForConditionalGeneration,
+ ... )
+
+ >>> # Initializing text encoder, audio encoder, and decoder model configurations
+ >>> text_encoder_config = T5Config()
+ >>> audio_encoder_config = EncodecConfig()
+ >>> decoder_config = MusicgenDecoderConfig()
+
+ >>> configuration = MusicgenConfig.from_sub_models_config(
+ ... text_encoder_config, audio_encoder_config, decoder_config
+ ... )
+
+ >>> # Initializing a MusicgenForConditionalGeneration (with random weights) from the facebook/musicgen-small style configuration
+ >>> model = MusicgenForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ >>> config_text_encoder = model.config.text_encoder
+ >>> config_audio_encoder = model.config.audio_encoder
+ >>> config_decoder = model.config.decoder
+
+ >>> # Saving the model, including its configuration
+ >>> model.save_pretrained("musicgen-model")
+
+ >>> # loading model and config from pretrained folder
+ >>> musicgen_config = MusicgenConfig.from_pretrained("musicgen-model")
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("musicgen-model", config=musicgen_config)
+ ```"""
+
+ model_type = "musicgen"
+ sub_configs = {
+ "text_encoder": AutoConfig,
+ "audio_encoder": AutoConfig,
+ "decoder": MusicgenDecoderConfig,
+ }
+ is_composition = True
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
+ raise ValueError("Config has to be initialized with text_encoder, audio_encoder and decoder config")
+
+ text_encoder_config = kwargs.pop("text_encoder")
+ text_encoder_model_type = text_encoder_config.pop("model_type")
+
+ audio_encoder_config = kwargs.pop("audio_encoder")
+ audio_encoder_model_type = audio_encoder_config.pop("model_type")
+
+ decoder_config = kwargs.pop("decoder")
+
+ self.text_encoder = AutoConfig.for_model(text_encoder_model_type, **text_encoder_config)
+ self.audio_encoder = AutoConfig.for_model(audio_encoder_model_type, **audio_encoder_config)
+ self.decoder = MusicgenDecoderConfig(**decoder_config)
+ self.is_encoder_decoder = True
+
+ @classmethod
+ def from_sub_models_config(
+ cls,
+ text_encoder_config: PretrainedConfig,
+ audio_encoder_config: PretrainedConfig,
+ decoder_config: MusicgenDecoderConfig,
+ **kwargs,
+ ):
+ r"""
+ Instantiate a [`MusicgenConfig`] (or a derived class) from text encoder, audio encoder and decoder
+ configurations.
+
+ Returns:
+ [`MusicgenConfig`]: An instance of a configuration object
+ """
+
+ return cls(
+ text_encoder=text_encoder_config.to_dict(),
+ audio_encoder=audio_encoder_config.to_dict(),
+ decoder=decoder_config.to_dict(),
+ **kwargs,
+ )
+
+ @property
+ # This is a property because you might want to change the codec model on the fly
+ def sampling_rate(self):
+ return self.audio_encoder.sampling_rate
+
+
+__all__ = ["MusicgenConfig", "MusicgenDecoderConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py b/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..151c4e89f3bff928a231b3bb6b78b9c874cb29ea
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py
@@ -0,0 +1,2741 @@
+# 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 Musicgen model."""
+
+import copy
+import inspect
+import math
+import random
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from ...activations import ACT2FN
+from ...generation import (
+ ClassifierFreeGuidanceLogitsProcessor,
+ GenerationConfig,
+ GenerationMixin,
+ GenerationMode,
+ LogitsProcessorList,
+ StoppingCriteriaList,
+)
+from ...modeling_attn_mask_utils import (
+ _prepare_4d_attention_mask,
+ _prepare_4d_attention_mask_for_sdpa,
+ _prepare_4d_causal_attention_mask,
+ _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ ModelOutput,
+ 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
+from .configuration_musicgen import MusicgenConfig, MusicgenDecoderConfig
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+if TYPE_CHECKING:
+ from ...generation.streamers import BaseStreamer
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MusicgenConfig"
+_CHECKPOINT_FOR_DOC = "facebook/musicgen-small"
+
+
+@dataclass
+class MusicgenUnconditionalInput(ModelOutput):
+ """
+ Args:
+ encoder_outputs (`Tuple[torch.FloatTensor]` of length 1, with tensor shape `(batch_size, sequence_length, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the text encoder model.
+ attention_mask (`torch.LongTensor`) of shape `(batch_size, sequence_length)`, *optional*):
+ Encoder attention 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**.
+ guidance_scale (`float`, *optional*):
+ Guidance scale for classifier free guidance, setting the balance between the conditional logits (predicted
+ from the prompts) and the unconditional logits (predicted without prompts).
+ """
+
+ encoder_outputs: Tuple[torch.FloatTensor] = None
+ attention_mask: Optional[torch.LongTensor] = None
+ guidance_scale: Optional[float] = None
+
+
+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.
+ """
+ # transpose to get (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.transpose(1, 2)
+ 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
+
+
+class MusicgenSinusoidalPositionalEmbedding(nn.Module):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, num_positions: int, embedding_dim: int):
+ super().__init__()
+ self.embedding_dim = embedding_dim
+ self.make_weights(num_positions, embedding_dim)
+
+ def make_weights(self, num_embeddings: int, embedding_dim: int):
+ emb_weights = self.get_embedding(num_embeddings, embedding_dim)
+ if hasattr(self, "weights"):
+ # in forward put the weights on the correct dtype and device of the param
+ emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
+
+ self.weights = nn.Parameter(emb_weights)
+ self.weights.requires_grad = False
+ self.weights.detach_()
+
+ @staticmethod
+ def get_embedding(num_embeddings: int, embedding_dim: int):
+ """
+ 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.cos(emb), torch.sin(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)
+ return emb.to(torch.get_default_dtype())
+
+ @torch.no_grad()
+ def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
+ bsz, codebooks, seq_len = input_ids.size()
+ # Create the position ids from the input token ids.
+ position_ids = (torch.arange(seq_len) + past_key_values_length).to(input_ids.device)
+ # expand embeddings if needed
+ if seq_len > self.weights.size(0):
+ self.make_weights(seq_len + self.offset, self.embedding_dim)
+ return self.weights.index_select(0, position_ids.view(-1)).detach()
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->Musicgen
+class MusicgenAttention(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[MusicgenConfig] = 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.bart.modeling_bart.BartFlashAttention2 with Bart->Musicgen
+class MusicgenFlashAttention2(MusicgenAttention):
+ """
+ Musicgen flash attention module. This module inherits from `MusicgenAttention` 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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ def _reshape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+
+ 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]]]:
+ # MusicgenFlashAttention2 attention does not support output_attentions
+ if output_attentions:
+ raise ValueError("MusicgenFlashAttention2 attention does not support output_attentions")
+
+ # 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, q_len, _ = hidden_states.size()
+
+ # get query proj
+ query_states = self._reshape(self.q_proj(hidden_states), -1, bsz)
+ # 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].transpose(1, 2)
+ value_states = past_key_value[1].transpose(1, 2)
+ elif is_cross_attention:
+ # cross_attentions
+ key_states = self._reshape(self.k_proj(key_value_states), -1, bsz)
+ value_states = self._reshape(self.v_proj(key_value_states), -1, bsz)
+ elif past_key_value is not None:
+ # reuse k, v, self_attention
+ key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)
+ key_states = torch.cat([past_key_value[0].transpose(1, 2), key_states], dim=1)
+ value_states = torch.cat([past_key_value[1].transpose(1, 2), value_states], dim=1)
+ else:
+ # self_attention
+ key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._reshape(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.transpose(1, 2), value_states.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]
+
+ # 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 = 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 = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=self.dropout if self.training else 0.0,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1)
+ attn_output = self.out_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+class MusicgenSdpaAttention(MusicgenAttention):
+ 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 output_attentions or layer_head_mask is not None:
+ # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.
+ logger.warning_once(
+ "MusicgenModel is using MusicgenSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_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.'
+ )
+ return super().forward(
+ hidden_states,
+ key_value_states=key_value_states,
+ past_key_value=past_key_value,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ output_attentions=output_attentions,
+ )
+
+ if (
+ attention_mask is not None
+ and (attention_mask.mean(dim=[1, 2, 3]) <= torch.finfo(attention_mask.dtype).min).any()
+ ):
+ logger.warning_once(
+ '`torch.nn.functional.scaled_dot_product_attention` does not support having an empty attention mask. Falling back to the manual attention implementation. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ "Note that this probably happens because `guidance_scale>1` or because you used `get_unconditional_inputs`. See https://github.com/huggingface/transformers/issues/31189 for more information."
+ )
+ return super().forward(
+ hidden_states,
+ key_value_states=key_value_states,
+ past_key_value=past_key_value,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ output_attentions=output_attentions,
+ )
+
+ # 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)
+ # 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)
+
+ query_states = self._shape(query_states, tgt_len, bsz)
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case tgt_len == 1.
+ is_causal = True if self.is_causal and attention_mask is None and tgt_len > 1 else False
+
+ # NOTE: SDPA with memory-efficient backend is currently (torch==2.1.2) bugged when using non-contiguous inputs and a custom attn_mask,
+ # but we are fine here as `_shape` do call `.contiguous()`. Reference: https://github.com/pytorch/pytorch/issues/112577
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attn_mask=attention_mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ is_causal=is_causal,
+ )
+
+ 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.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, None, past_key_value
+
+
+MUSICGEN_ATTENTION_CLASSES = {
+ "eager": MusicgenAttention,
+ "sdpa": MusicgenSdpaAttention,
+ "flash_attention_2": MusicgenFlashAttention2,
+}
+
+
+class MusicgenDecoderLayer(nn.Module):
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+
+ self.self_attn = MUSICGEN_ATTENTION_CLASSES[config._attn_implementation](
+ embed_dim=self.embed_dim,
+ num_heads=config.num_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ is_causal=True,
+ config=config,
+ )
+ 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 = MUSICGEN_ATTENTION_CLASSES[config._attn_implementation](
+ self.embed_dim,
+ config.num_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ config=config,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.ffn_dim, bias=False)
+ self.fc2 = nn.Linear(config.ffn_dim, self.embed_dim, bias=False)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ # Copied from transformers.models.mbart.modeling_mbart.MBartDecoderLayer.forward
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ layer_head_mask: Optional[torch.Tensor] = None,
+ cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = True,
+ ) -> torch.Tensor:
+ """
+ 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.
+ 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, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+ `(encoder_attention_heads,)`.
+ cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of
+ size `(decoder_attention_heads,)`.
+ past_key_value (`Tuple(torch.FloatTensor)`): 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.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # 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
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_value=self_attn_past_key_value,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ output_attentions=output_attentions,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Cross-Attention Block
+ cross_attn_present_key_value = None
+ cross_attn_weights = None
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
+ cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+ hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ layer_head_mask=cross_attn_layer_head_mask,
+ past_key_value=cross_attn_past_key_value,
+ output_attentions=output_attentions,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # add cross-attn to positions 3,4 of present_key_value tuple
+ present_key_value = present_key_value + cross_attn_present_key_value
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(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
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights, cross_attn_weights)
+
+ if use_cache:
+ outputs += (present_key_value,)
+
+ return outputs
+
+
+class MusicgenPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MusicgenDecoderConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MusicgenDecoderLayer", "MusicgenAttention"]
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_factor
+ 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_()
+ 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, MusicgenSinusoidalPositionalEmbedding):
+ weights = module.get_embedding(*module.weights.shape)
+ weights = nn.Parameter(weights, requires_grad=False)
+ weights.detach_()
+ module.weights = weights
+
+
+MUSICGEN_START_DOCSTRING = r"""
+
+ The Musicgen model was proposed in [Simple and Controllable Music Generation](https://arxiv.org/abs/2306.05284) by
+ Jade Copet, Felix Kreuk, Itai Gat, Tal Remez, David Kant, Gabriel Synnaeve, Yossi Adi, Alexandre Défossez. It is an
+ encoder decoder transformer trained on the task of conditional music generation
+
+ 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 ([`MusicgenConfig`]): 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.
+"""
+
+MUSICGEN_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)
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+
+
+ The `decoder_input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `decoder_input_ids`.
+
+
+
+ decoder_attention_mask (`torch.LongTensor` 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(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.
+ 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.
+ 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`.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *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]`
+ 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.
+"""
+
+MUSICGEN_DECODER_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+
+
+
+ The `input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `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)
+ 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 tokens indices of encoder input_ids. 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.Tensor` of shape `(decoder_layers, decoder_attention_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 `(decoder_layers, decoder_attention_heads)`, *optional*):
+ Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing
+ cross-attention on hidden heads. Mask values selected in `[0, 1]`:
+
+ - 1 indicates the head is **not masked**,
+ - 0 indicates the head is **masked**.
+
+ 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.
+ 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 MusicgenDecoder(MusicgenPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MusicgenDecoderLayer`]
+ """
+
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.layerdrop
+ self.max_target_positions = config.max_position_embeddings
+ self.d_model = config.hidden_size
+ self.num_codebooks = config.num_codebooks
+ self.embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
+
+ embed_dim = config.vocab_size + 1
+ self.embed_tokens = nn.ModuleList(
+ [nn.Embedding(embed_dim, config.hidden_size) for _ in range(config.num_codebooks)]
+ )
+
+ self.embed_positions = MusicgenSinusoidalPositionalEmbedding(
+ config.max_position_embeddings,
+ config.hidden_size,
+ )
+
+ self.layers = nn.ModuleList([MusicgenDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.layer_norm = nn.LayerNorm(config.hidden_size)
+ self.attn_implementation = config._attn_implementation
+
+ 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(MUSICGEN_DECODER_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ cross_attn_head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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, 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
+
+ # 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:
+ # (bsz * codebooks, seq_len) -> (bsz, codebooks, seq_len)
+ input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input.shape
+ input_shape = (bsz, seq_len)
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ input = inputs_embeds[:, :, -1:]
+ else:
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+ # 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 inputs_embeds is None:
+ inputs_embeds = sum([self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks)])
+
+ if self.attn_implementation == "flash_attention_2":
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ elif self.attn_implementation == "sdpa" and head_mask is None and not output_attentions:
+ # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+ # the manual implementation that requires a 4D causal mask in all cases.
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+ attention_mask,
+ input_shape,
+ inputs_embeds,
+ past_key_values_length,
+ )
+ else:
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask, input_shape, inputs_embeds, past_key_values_length
+ )
+
+ # expand encoder attention mask
+ if encoder_hidden_states is not None and encoder_attention_mask is not None:
+ if self.attn_implementation == "flash_attention_2":
+ encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None
+ elif self.attn_implementation == "sdpa" and cross_attn_head_mask is None and not output_attentions:
+ # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on
+ # the manual implementation that requires a 4D causal mask in all cases.
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+ encoder_attention_mask,
+ inputs_embeds.dtype,
+ tgt_len=input_shape[-1],
+ )
+ else:
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ encoder_attention_mask = _prepare_4d_attention_mask(
+ encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+ )
+
+ # embed positions
+ positions = self.embed_positions(input, past_key_values_length)
+
+ hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ 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
+
+ # 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
+ next_decoder_cache = () if use_cache else None
+
+ # check if head_mask/cross_attn_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:
+ if attn_mask.size()[0] != len(self.layers):
+ raise ValueError(
+ f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
+ f" {attn_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:
+ all_hidden_states += (hidden_states,)
+ dropout_probability = random.uniform(0, 1)
+ if self.training and (dropout_probability < self.layerdrop):
+ continue
+
+ 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.forward,
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ head_mask[idx] if head_mask is not None else None,
+ cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
+ None,
+ output_attentions,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_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
+ ),
+ 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[3 if output_attentions else 1],)
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if encoder_hidden_states is not None:
+ all_cross_attentions += (layer_outputs[2],)
+
+ hidden_states = self.layer_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, all_cross_attentions]
+ if v is not None
+ )
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+@add_start_docstrings(
+ "The bare Musicgen decoder model outputting raw hidden-states without any specific head on top.",
+ MUSICGEN_START_DOCSTRING,
+)
+class MusicgenModel(MusicgenPreTrainedModel):
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+ self.decoder = MusicgenDecoder(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.decoder.embed_tokens = value
+
+ def get_decoder(self):
+ return self.decoder
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_DECODER_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ cross_attn_head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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, 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
+
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
+ decoder_outputs = self.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ head_mask=head_mask,
+ cross_attn_head_mask=cross_attn_head_mask,
+ 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,
+ )
+
+ if not return_dict:
+ return decoder_outputs
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ hidden_states=decoder_outputs.hidden_states,
+ attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ )
+
+
+@add_start_docstrings(
+ "The MusicGen decoder model with a language modelling head on top.",
+ MUSICGEN_START_DOCSTRING,
+)
+class MusicgenForCausalLM(MusicgenPreTrainedModel, GenerationMixin):
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+
+ self.model = MusicgenModel(config)
+
+ self.num_codebooks = config.num_codebooks
+ self.lm_heads = nn.ModuleList(
+ [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_codebooks)]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_heads
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_heads = new_embeddings
+
+ def set_decoder(self, decoder):
+ self.model.decoder = decoder
+
+ def get_decoder(self):
+ return self.model.decoder
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_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.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ cross_attn_head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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,
+ **kwargs,
+ ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *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:
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if (labels is not None) and (input_ids is None and inputs_embeds is None):
+ input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.bos_token_id)
+
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ head_mask=head_mask,
+ cross_attn_head_mask=cross_attn_head_mask,
+ 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]
+
+ lm_logits = torch.stack([head(hidden_states) for head in self.lm_heads], dim=1)
+
+ loss = None
+ if labels is not None:
+ # since encoder hidden states have been concatenated to the decoder hidden states,
+ # we take the last timestamps corresponding to labels
+ logits = lm_logits[:, :, -labels.shape[1] :]
+
+ loss_fct = CrossEntropyLoss()
+ loss = torch.zeros([], device=self.device)
+
+ # per codebook cross-entropy
+ # -100 labels are ignored
+ labels = labels.masked_fill(labels == self.config.pad_token_id, -100)
+
+ # per codebook cross-entropy
+ # ref: https://github.com/facebookresearch/audiocraft/blob/69fea8b290ad1b4b40d28f92d1dfc0ab01dbab85/audiocraft/solvers/musicgen.py#L242-L243
+ for codebook in range(self.config.num_codebooks):
+ codebook_logits = logits[:, codebook].contiguous().view(-1, logits.shape[-1])
+ codebook_labels = labels[..., codebook].contiguous().view(-1)
+ loss += loss_fct(codebook_logits, codebook_labels)
+
+ loss = loss / self.config.num_codebooks
+
+ # (bsz, num_codebooks, seq_len, vocab_size) -> (bsz * num_codebooks, seq_len, vocab_size)
+ lm_logits = lm_logits.reshape(-1, *lm_logits.shape[2:])
+
+ if not return_dict:
+ output = (lm_logits,) + outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=lm_logits,
+ 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,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ head_mask=None,
+ cross_attn_head_mask=None,
+ past_key_values=None,
+ use_cache=True,
+ delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if delay_pattern_mask is None:
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ input_ids = self.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ input_ids = input_ids.repeat((2, 1))
+ if attention_mask is not None:
+ attention_mask = attention_mask.repeat((2, 1))
+
+ if past_key_values is not None:
+ input_ids = input_ids[:, -1:]
+
+ return {
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "encoder_hidden_states": encoder_hidden_states,
+ "encoder_attention_mask": encoder_attention_mask,
+ "head_mask": head_mask,
+ "cross_attn_head_mask": cross_attn_head_mask,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ }
+
+ def build_delay_pattern_mask(
+ self, input_ids: torch.LongTensor, pad_token_id: int, max_length: Optional[int] = None
+ ):
+ """Build a delayed pattern mask to the input_ids. Each codebook is offset by the previous codebook by
+ one, giving a delayed pattern mask at the start of sequence and end of sequence. Take the example where there
+ are 4 codebooks and a max sequence length of 8, we have the delayed pattern mask of shape `(codebooks,
+ seq_len)`:
+ - [P, -1, -1, -1, -1, P, P, P]
+ - [P, P, -1, -1, -1, -1, P, P]
+ - [P, P, P, -1, -1, -1, -1, P]
+ - [P, P, P, P, -1, -1, -1, -1]
+ where P is the special padding token id and -1 indicates that the token is valid for prediction. If we include
+ a prompt (decoder input ids), the -1 positions indicate where new tokens should be predicted. Otherwise, the
+ mask is set to the value in the prompt:
+ - [P, a, b, -1, -1, P, P, P]
+ - [P, P, c, d, -1, -1, P, P]
+ - [P, P, P, e, f, -1, -1, P]
+ - [P, P, P, P, g, h, -1, -1]
+ where a-h indicate the input prompt (decoder input ids) that are offset by 1. Now, we only override the -1
+ tokens in our prediction.
+ """
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input_ids.shape
+
+ max_length = max_length if max_length is not None else self.generation_config.max_length
+ input_ids_shifted = (
+ torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1
+ )
+
+ channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks
+ # we only apply the mask if we have a large enough seq len - otherwise we return as is
+ if max_length < 2 * channel_codebooks - 1:
+ return input_ids.reshape(bsz * num_codebooks, -1), input_ids_shifted.reshape(bsz * num_codebooks, -1)
+
+ # fill the shifted ids with the prompt entries, offset by the codebook idx
+ for codebook in range(channel_codebooks):
+ if self.config.audio_channels == 1:
+ # mono channel - loop over the codebooks one-by-one
+ input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook]
+ else:
+ # left/right channels are interleaved in the generated codebooks, so handle one then the other
+ input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook]
+ input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1]
+
+ # construct a pattern mask that indicates the positions of padding tokens for each codebook
+ # first fill the upper triangular part (the EOS padding)
+ delay_pattern = torch.triu(
+ torch.ones((channel_codebooks, max_length), dtype=torch.bool), diagonal=max_length - channel_codebooks + 1
+ )
+ # then fill the lower triangular part (the BOS padding)
+ delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.bool))
+
+ if self.config.audio_channels == 2:
+ # for left/right channel we need to duplicate every row of the pattern mask in an interleaved fashion
+ delay_pattern = delay_pattern.repeat_interleave(2, dim=0)
+
+ mask = ~delay_pattern.to(input_ids.device)
+ input_ids = mask * input_ids_shifted + ~mask * pad_token_id
+
+ # find the first position to start generating - this is the first place we have the -1 token
+ # and will always be in the first codebook (since it has no codebook offset)
+ first_codebook_ids = input_ids[:, 0, :]
+ start_ids = (first_codebook_ids == -1).nonzero()[:, 1]
+ if len(start_ids) > 0:
+ first_start_id = min(start_ids)
+ else:
+ # we have no tokens that need to be filled - return entire matrix of input ids
+ first_start_id = seq_len
+
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ pattern_mask = input_ids.reshape(bsz * num_codebooks, -1)
+ input_ids = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1)
+ return input_ids, pattern_mask
+
+ @staticmethod
+ def apply_delay_pattern_mask(input_ids, decoder_pad_token_mask):
+ """Apply a delay pattern mask to the decoder input ids, only preserving predictions where
+ the mask is set to -1, and otherwise setting to the value detailed in the mask."""
+ seq_len = input_ids.shape[-1]
+ decoder_pad_token_mask = decoder_pad_token_mask[..., :seq_len]
+ input_ids = torch.where(decoder_pad_token_mask == -1, input_ids, decoder_pad_token_mask)
+ return input_ids
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ synced_gpus: Optional[bool] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ if generation_config is None:
+ generation_config = self.generation_config
+
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+ generation_config.validate()
+ self._validate_model_kwargs(model_kwargs.copy())
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs`
+ input_ids, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = input_ids.shape[0] // self.num_codebooks
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=input_ids.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ input_ids, generation_config, model_kwargs
+ )
+
+ # 5. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=input_ids,
+ input_ids_length=input_ids_length,
+ )
+
+ # 6. Prepare `input_ids` which will be used for auto-regressive generation
+ # Build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # stash the delay mask so that we don't have to recompute it in each forward pass
+ model_kwargs["delay_pattern_mask"] = delay_pattern_mask
+
+ # 7. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 8. prepare batched CFG externally (to enable coexistance with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 9. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=input_ids,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ **model_kwargs,
+ )
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=synced_gpus,
+ streamer=streamer,
+ **model_kwargs,
+ )
+
+ else:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.apply_delay_pattern_mask(output_ids, model_kwargs["delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.num_codebooks, -1
+ )
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_ids
+ return outputs
+ else:
+ return output_ids
+
+
+@add_start_docstrings(
+ "The composite MusicGen model with a text encoder, audio encoder and Musicgen decoder, "
+ "for music generation tasks with one or both of text and audio prompts.",
+ MUSICGEN_START_DOCSTRING,
+)
+class MusicgenForConditionalGeneration(PreTrainedModel, GenerationMixin):
+ config_class = MusicgenConfig
+ base_model_prefix = "encoder_decoder"
+ main_input_name = "input_ids"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ def __init__(
+ self,
+ config: Optional[MusicgenConfig] = None,
+ text_encoder: Optional[PreTrainedModel] = None,
+ audio_encoder: Optional[PreTrainedModel] = None,
+ decoder: Optional[MusicgenForCausalLM] = None,
+ ):
+ if config is None and (text_encoder is None or audio_encoder is None or decoder is None):
+ raise ValueError(
+ "Either a configuration has to be provided, or all three of text encoder, audio encoder and MusicGen decoder."
+ )
+ if config is None:
+ config = MusicgenConfig.from_sub_models_config(text_encoder.config, audio_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.text_encoder.hidden_size:
+ raise ValueError(
+ "If `cross_attention_hidden_size` is specified in the MusicGen decoder's configuration, it has to be equal"
+ f" to the text encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
+ f" `config.decoder.cross_attention_hidden_size` and {config.text_encoder.hidden_size} for"
+ " `config.text_encoder.hidden_size`."
+ )
+
+ # initialize with config
+ super().__init__(config)
+
+ if text_encoder is None:
+ from ..auto.modeling_auto import AutoModelForTextEncoding
+
+ text_encoder = AutoModelForTextEncoding.from_config(config.text_encoder)
+
+ if audio_encoder is None:
+ from ..auto.modeling_auto import AutoModel
+
+ audio_encoder = AutoModel.from_config(config.audio_encoder)
+
+ if decoder is None:
+ decoder = MusicgenForCausalLM._from_config(config.decoder)
+
+ self.text_encoder = text_encoder
+ self.audio_encoder = audio_encoder
+ self.decoder = decoder
+
+ if self.text_encoder.config.to_dict() != self.config.text_encoder.to_dict():
+ logger.warning(
+ f"Config of the text_encoder: {self.text_encoder.__class__} is overwritten by shared text_encoder config:"
+ f" {self.config.text_encoder}"
+ )
+ if self.audio_encoder.config.to_dict() != self.config.audio_encoder.to_dict():
+ logger.warning(
+ f"Config of the audio_encoder: {self.audio_encoder.__class__} is overwritten by shared audio_encoder config:"
+ f" {self.config.audio_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.config.text_encoder._attn_implementation = self.text_encoder.config._attn_implementation
+ self.config.audio_encoder._attn_implementation = self.audio_encoder.config._attn_implementation
+ self.config.decoder._attn_implementation = self.decoder.config._attn_implementation
+ self.text_encoder.config = self.config.text_encoder
+ self.audio_encoder.config = self.config.audio_encoder
+ self.decoder.config = self.config.decoder
+
+ # text encoder outputs might need to be projected to different dimension for decoder
+ if (
+ self.text_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.text_encoder.config.hidden_size, self.decoder.config.hidden_size)
+
+ if self.text_encoder.get_output_embeddings() is not None:
+ raise ValueError(
+ f"The encoder {self.text_encoder} should not have a LM Head. Please use a model without and 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 text encoder, decoder weights if config set accordingly
+ self.tie_weights()
+
+ def tie_weights(self):
+ # tie text encoder & decoder if needed
+ if self.config.tie_encoder_decoder:
+ # tie text encoder and decoder base model
+ decoder_base_model_prefix = self.decoder.base_model_prefix
+ tied_weights = self._tie_encoder_decoder_weights(
+ self.text_encoder,
+ self.decoder._modules[decoder_base_model_prefix],
+ self.decoder.base_model_prefix,
+ "text_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_audio_encoder(self):
+ return self.audio_encoder
+
+ def get_text_encoder(self):
+ return self.text_encoder
+
+ def get_encoder(self):
+ # get the text encoder to compute the encoder hidden-states for generation
+ return self.get_text_encoder()
+
+ def get_decoder(self):
+ return self.decoder
+
+ def get_input_embeddings(self):
+ return self.text_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_sub_models_pretrained(
+ cls,
+ text_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ audio_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ decoder_pretrained_model_name_or_path: Optional[str] = None,
+ *model_args,
+ **kwargs,
+ ) -> PreTrainedModel:
+ r"""
+ Instantiate a text encoder, an audio encoder, and a MusicGen decoder from one, two or three 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:
+ text_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text 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/`.
+
+ audio_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the audio 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/`.
+
+ 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/`.
+
+ 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 text encoder configuration, use the prefix *text_encoder_* for each configuration
+ parameter.
+ - To update the audio encoder configuration, use the prefix *audio_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 MusicgenForConditionalGeneration
+
+ >>> # initialize a musicgen model from a t5 text encoder, encodec audio encoder, and musicgen decoder
+ >>> model = MusicgenForConditionalGeneration.from_sub_models_pretrained(
+ ... text_encoder_pretrained_model_name_or_path="google-t5/t5-base",
+ ... audio_encoder_pretrained_model_name_or_path="facebook/encodec_24khz",
+ ... decoder_pretrained_model_name_or_path="facebook/musicgen-small",
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./musicgen-ft")
+ >>> # load fine-tuned model
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("./musicgen-ft")
+ ```"""
+
+ kwargs_text_encoder = {
+ argument[len("text_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("text_encoder_")
+ }
+
+ kwargs_audio_encoder = {
+ argument[len("audio_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("audio_encoder_")
+ }
+
+ kwargs_decoder = {
+ argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+ }
+
+ # remove text encoder, audio encoder and decoder kwargs from kwargs
+ for key in kwargs_text_encoder.keys():
+ del kwargs["text_encoder_" + key]
+ for key in kwargs_audio_encoder.keys():
+ del kwargs["audio_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.
+ text_encoder = kwargs_text_encoder.pop("model", None)
+ if text_encoder is None:
+ if text_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `text_encoder_model` is not defined as an argument, a `text_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_text_encoder:
+ encoder_config, kwargs_text_encoder = AutoConfig.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, **kwargs_text_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {text_encoder_pretrained_model_name_or_path} as a text_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_text_encoder["config"] = encoder_config
+
+ text_encoder = AutoModel.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, *model_args, **kwargs_text_encoder
+ )
+
+ audio_encoder = kwargs_audio_encoder.pop("model", None)
+ if audio_encoder is None:
+ if audio_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `audio_encoder_model` is not defined as an argument, an `audio_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_audio_encoder:
+ encoder_config, kwargs_audio_encoder = AutoConfig.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, **kwargs_audio_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {audio_encoder_pretrained_model_name_or_path} as an audio_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_audio_encoder["config"] = encoder_config
+
+ audio_encoder = AutoModel.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, *model_args, **kwargs_audio_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 isinstance(decoder_config, MusicgenConfig):
+ decoder_config = decoder_config.decoder
+
+ 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_sub_models_pretrained(...)` are set to `True` or do not pass a "
+ "`decoder_config` to `.from_sub_models_pretrained(...)`"
+ )
+
+ decoder = MusicgenForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
+
+ # instantiate config with corresponding kwargs
+ config = MusicgenConfig.from_sub_models_config(
+ text_encoder.config, audio_encoder.config, decoder.config, **kwargs
+ )
+ return cls(text_encoder=text_encoder, audio_encoder=audio_encoder, decoder=decoder, config=config)
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_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.BoolTensor] = None,
+ input_values: Optional[torch.FloatTensor] = None,
+ padding_mask: Optional[torch.BoolTensor] = 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 AutoProcessor, MusicgenForConditionalGeneration
+ >>> import torch
+
+ >>> processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+ >>> inputs = processor(
+ ... text=["80s pop track with bassy drums and synth", "90s rock song with loud guitars and heavy drums"],
+ ... padding=True,
+ ... return_tensors="pt",
+ ... )
+
+ >>> pad_token_id = model.generation_config.pad_token_id
+ >>> decoder_input_ids = (
+ ... torch.ones((inputs.input_ids.shape[0] * model.decoder.num_codebooks, 1), dtype=torch.long)
+ ... * pad_token_id
+ ... )
+
+ >>> logits = model(**inputs, decoder_input_ids=decoder_input_ids).logits
+ >>> logits.shape # (bsz * num_codebooks, tgt_len, vocab_size)
+ torch.Size([8, 1, 2048])
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ kwargs_text_encoder = {
+ argument[len("text_encoder_")]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("text_encoder_")
+ }
+
+ kwargs_audio_encoder = {
+ argument[len("audio_encoder_")]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("audio_encoder_")
+ }
+
+ kwargs_decoder = {
+ argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+ }
+
+ if encoder_outputs is None:
+ encoder_outputs = self.text_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_text_encoder,
+ )
+ elif isinstance(encoder_outputs, tuple):
+ encoder_outputs = BaseModelOutput(*encoder_outputs)
+
+ encoder_hidden_states = encoder_outputs[0]
+
+ # optionally project encoder_hidden_states
+ if (
+ self.text_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 attention_mask is not None:
+ encoder_hidden_states = encoder_hidden_states * attention_mask[..., None]
+
+ 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.decoder.pad_token_id, self.config.decoder.decoder_start_token_id
+ )
+
+ elif decoder_input_ids is None and decoder_inputs_embeds is None:
+ audio_encoder_outputs = self.audio_encoder(
+ input_values=input_values,
+ padding_mask=padding_mask,
+ **kwargs_audio_encoder,
+ )
+ audio_codes = audio_encoder_outputs.audio_codes
+ frames, bsz, codebooks, seq_len = audio_codes.shape
+ if frames != 1:
+ raise ValueError(
+ f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
+ "disabled by setting `chunk_length=None` in the audio encoder."
+ )
+
+ if self.config.decoder.audio_channels == 2 and audio_codes.shape[2] == self.decoder.num_codebooks // 2:
+ # mono input through encodec that we convert to stereo
+ audio_codes = audio_codes.repeat_interleave(2, dim=2)
+
+ decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)
+
+ # 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,
+ labels=labels,
+ **kwargs_decoder,
+ )
+
+ if not return_dict:
+ return decoder_outputs + encoder_outputs
+
+ return Seq2SeqLMOutput(
+ loss=decoder_outputs.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,
+ decoder_input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ head_mask=None,
+ decoder_attention_mask=None,
+ decoder_head_mask=None,
+ cross_attn_head_mask=None,
+ use_cache=None,
+ encoder_outputs=None,
+ decoder_delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if decoder_delay_pattern_mask is None:
+ decoder_input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ decoder_input_ids,
+ self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ decoder_input_ids = self.decoder.apply_delay_pattern_mask(decoder_input_ids, decoder_delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ decoder_input_ids = decoder_input_ids.repeat((2, 1))
+ if decoder_attention_mask is not None:
+ decoder_attention_mask = decoder_attention_mask.repeat((2, 1))
+
+ 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 decoder_input_ids.shape[1] > past_length:
+ remove_prefix_length = past_length
+ else:
+ # Default to old behavior: keep only final ID
+ remove_prefix_length = decoder_input_ids.shape[1] - 1
+
+ decoder_input_ids = decoder_input_ids[:, remove_prefix_length:]
+
+ 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,
+ "decoder_attention_mask": decoder_attention_mask,
+ "head_mask": head_mask,
+ "decoder_head_mask": decoder_head_mask,
+ "cross_attn_head_mask": cross_attn_head_mask,
+ "use_cache": use_cache,
+ }
+
+ def _prepare_decoder_input_ids_for_generation(
+ self,
+ batch_size: int,
+ model_input_name: str,
+ model_kwargs: Dict[str, torch.Tensor],
+ decoder_start_token_id: Optional[int] = None,
+ bos_token_id: Optional[int] = None,
+ device: torch.device = None,
+ ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]:
+ """Prepares `decoder_input_ids` for generation with encoder-decoder models"""
+
+ # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming,
+ # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input.
+ if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
+ decoder_input_ids = model_kwargs.pop("decoder_input_ids")
+ elif "input_ids" in model_kwargs and model_input_name != "input_ids":
+ decoder_input_ids = model_kwargs.pop("input_ids")
+ else:
+ decoder_input_ids = None
+
+ # 2. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that.
+ decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)
+ if device is None:
+ device = self.device
+ decoder_input_ids_start = (
+ torch.ones((batch_size * self.decoder.num_codebooks, 1), dtype=torch.long, device=device)
+ * decoder_start_token_id
+ )
+
+ # no user input -> use decoder_start_token_id as decoder_input_ids
+ if decoder_input_ids is None:
+ decoder_input_ids = decoder_input_ids_start
+
+ # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust
+ # decoder_attention_mask if provided)
+ elif (decoder_input_ids[..., 0] != decoder_start_token_id).all().item():
+ decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
+ if "decoder_attention_mask" in model_kwargs:
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+ decoder_attention_mask = torch.cat(
+ (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),
+ dim=-1,
+ )
+ model_kwargs["decoder_attention_mask"] = decoder_attention_mask
+
+ return decoder_input_ids, model_kwargs
+
+ def _prepare_text_encoder_kwargs_for_generation(
+ self,
+ inputs_tensor: torch.Tensor,
+ model_kwargs,
+ model_input_name: Optional[str],
+ generation_config: GenerationConfig,
+ ) -> Dict[str, Any]:
+ # 1. get text encoder
+ encoder = self.get_text_encoder()
+ # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+ # as the inputs.
+ if hasattr(encoder, "_hf_hook"):
+ encoder._hf_hook.io_same_device = True
+
+ # 2. Prepare encoder args and encoder kwargs from model kwargs.
+ irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+ encoder_kwargs = {
+ argument: value
+ for argument, value in model_kwargs.items()
+ if not any(argument.startswith(p) for p in irrelevant_prefix)
+ }
+ encoder_signature = set(inspect.signature(encoder.forward).parameters)
+ encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+ if not encoder_accepts_wildcard:
+ encoder_kwargs = {
+ argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+ }
+ encoder_kwargs["output_attentions"] = generation_config.output_attentions
+ encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+ guidance_scale = generation_config.guidance_scale
+
+ # 3. make sure that encoder returns `ModelOutput`
+ model_input_name = model_input_name if model_input_name is not None else self.text_encoder.main_input_name
+ encoder_kwargs["return_dict"] = True
+ encoder_kwargs[model_input_name] = inputs_tensor
+ last_hidden_state = encoder(**encoder_kwargs).last_hidden_state
+
+ # for classifier free guidance we need to add a 'null' input to our encoder hidden states
+ if guidance_scale is not None and guidance_scale > 1:
+ last_hidden_state = torch.concatenate([last_hidden_state, torch.zeros_like(last_hidden_state)], dim=0)
+ if "attention_mask" in model_kwargs:
+ model_kwargs["attention_mask"] = torch.concatenate(
+ [model_kwargs["attention_mask"], torch.zeros_like(model_kwargs["attention_mask"])], dim=0
+ )
+
+ model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=last_hidden_state)
+
+ return model_kwargs
+
+ def _prepare_audio_encoder_kwargs_for_generation(
+ self, input_values, model_kwargs, model_input_name: Optional[str] = None
+ ):
+ # 1. get audio encoder
+ encoder = self.get_audio_encoder()
+ # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+ # as the inputs.
+ if hasattr(encoder, "_hf_hook"):
+ encoder._hf_hook.io_same_device = True
+
+ # 2. Prepare encoder args and encoder kwargs from model kwargs.
+ irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+ encoder_kwargs = {
+ argument: value
+ for argument, value in model_kwargs.items()
+ if not any(argument.startswith(p) for p in irrelevant_prefix)
+ }
+ encoder_signature = set(inspect.signature(encoder.forward).parameters)
+ encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+ if not encoder_accepts_wildcard:
+ encoder_kwargs = {
+ argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+ }
+
+ # 3. make sure that encoder returns `ModelOutput`
+ model_input_name = model_input_name if model_input_name is not None else self.audio_encoder.main_input_name
+ encoder_kwargs["return_dict"] = True
+
+ if self.decoder.config.audio_channels == 1:
+ encoder_kwargs[model_input_name] = input_values
+ audio_encoder_outputs = encoder.encode(**encoder_kwargs)
+ audio_codes = audio_encoder_outputs.audio_codes
+ audio_scales = audio_encoder_outputs.audio_scales
+
+ frames, bsz, codebooks, seq_len = audio_codes.shape
+
+ else:
+ if input_values.shape[1] != 2:
+ raise ValueError(
+ f"Expected stereo audio (2-channels) but example has {input_values.shape[1]} channel."
+ )
+
+ encoder_kwargs[model_input_name] = input_values[:, :1, :]
+ audio_encoder_outputs_left = encoder.encode(**encoder_kwargs)
+ audio_codes_left = audio_encoder_outputs_left.audio_codes
+ audio_scales_left = audio_encoder_outputs_left.audio_scales
+
+ encoder_kwargs[model_input_name] = input_values[:, 1:, :]
+ audio_encoder_outputs_right = encoder.encode(**encoder_kwargs)
+ audio_codes_right = audio_encoder_outputs_right.audio_codes
+ audio_scales_right = audio_encoder_outputs_right.audio_scales
+
+ frames, bsz, codebooks, seq_len = audio_codes_left.shape
+ # copy alternating left/right channel codes into stereo codebook
+ audio_codes = audio_codes_left.new_ones((frames, bsz, 2 * codebooks, seq_len))
+
+ audio_codes[:, :, ::2, :] = audio_codes_left
+ audio_codes[:, :, 1::2, :] = audio_codes_right
+
+ if audio_scales_left != [None] or audio_scales_right != [None]:
+ audio_scales = torch.stack([audio_scales_left, audio_scales_right], dim=1)
+ else:
+ audio_scales = [None] * bsz
+
+ if frames != 1:
+ raise ValueError(
+ f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
+ "disabled by setting `chunk_length=None` in the audio encoder."
+ )
+
+ decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)
+
+ model_kwargs["decoder_input_ids"] = decoder_input_ids
+ model_kwargs["audio_scales"] = audio_scales
+ return model_kwargs
+
+ def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+ return shift_tokens_right(labels, self.config.decoder.pad_token_id, self.config.decoder.bos_token_id)
+
+ 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 freeze_audio_encoder(self):
+ """
+ Freeze the audio encoder weights.
+ """
+ for param in self.audio_encoder.parameters():
+ param.requires_grad = False
+ self.audio_encoder._requires_grad = False
+
+ def freeze_text_encoder(self):
+ """
+ Freeze the text encoder weights.
+ """
+ for param in self.text_encoder.parameters():
+ param.requires_grad = False
+ self.text_encoder._requires_grad = False
+
+ def _maybe_initialize_input_ids_for_generation(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ bos_token_id: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+ ) -> torch.LongTensor:
+ """Initializes input ids for generation, if necessary."""
+ if inputs is not None:
+ return inputs
+
+ encoder_outputs = model_kwargs.get("encoder_outputs")
+ if encoder_outputs is not None:
+ # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding
+ shape = encoder_outputs[0].size()[:-1]
+ return torch.ones(shape, dtype=torch.long, device=self.device) * -100
+
+ if bos_token_id is None:
+ raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")
+
+ # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with
+ # soft-prompting or in multimodal implementations built on top of decoder-only language models.
+ batch_size = 1
+ for value in model_kwargs.values():
+ if isinstance(value, torch.Tensor):
+ batch_size = value.shape[0]
+ break
+ return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id
+
+ def _get_decoder_start_token_id(
+ self, decoder_start_token_id: Union[int, List[int]] = None, bos_token_id: Optional[int] = None
+ ) -> int:
+ decoder_start_token_id = (
+ decoder_start_token_id
+ if decoder_start_token_id is not None
+ else self.generation_config.decoder_start_token_id
+ )
+ bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
+
+ if decoder_start_token_id is not None:
+ return decoder_start_token_id
+ elif bos_token_id is not None:
+ return bos_token_id
+ raise ValueError(
+ "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ synced_gpus: Optional[bool] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ if generation_config is None:
+ generation_config = self.generation_config
+
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+ generation_config.validate()
+ self._validate_model_kwargs(model_kwargs.copy())
+
+ if model_kwargs.get("encoder_outputs") is not None and type(model_kwargs["encoder_outputs"]) is tuple:
+ # wrap the unconditional outputs as a BaseModelOutput for compatibility with the rest of generate
+ model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=model_kwargs["encoder_outputs"][0])
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs
+ inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = inputs_tensor.shape[0]
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=inputs_tensor.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ inputs_tensor, generation_config, model_kwargs
+ )
+
+ if "encoder_outputs" not in model_kwargs:
+ # encoder_outputs are created and added to `model_kwargs`
+ model_kwargs = self._prepare_text_encoder_kwargs_for_generation(
+ inputs_tensor, model_kwargs, model_input_name, generation_config
+ )
+
+ if "decoder_input_ids" not in model_kwargs and "input_values" in model_kwargs:
+ model_kwargs = self._prepare_audio_encoder_kwargs_for_generation(
+ model_kwargs["input_values"],
+ model_kwargs,
+ )
+
+ # 5. Prepare `input_ids` which will be used for auto-regressive generation
+ input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
+ batch_size=batch_size,
+ model_input_name=model_input_name,
+ model_kwargs=model_kwargs,
+ decoder_start_token_id=generation_config._decoder_start_token_tensor,
+ bos_token_id=generation_config._bos_token_tensor,
+ device=inputs_tensor.device,
+ )
+
+ # 6. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=inputs_tensor,
+ input_ids_length=input_ids_length,
+ )
+
+ # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+ input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+ # stash the delay mask so that we don't have to recompute in each forward pass
+ model_kwargs["decoder_delay_pattern_mask"] = decoder_delay_pattern_mask
+
+ # input_ids are ready to be placed on the streamer (if used)
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # 7. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 8. prepare batched CFG externally (to enable coexistance with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 9. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=inputs_tensor,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ is_encoder_decoder=self.config.is_encoder_decoder,
+ **model_kwargs,
+ )
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=synced_gpus,
+ streamer=streamer,
+ **model_kwargs,
+ )
+
+ else:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.decoder.apply_delay_pattern_mask(output_ids, model_kwargs["decoder_delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.decoder.num_codebooks, -1
+ )
+
+ # append the frame dimension back to the audio codes
+ output_ids = output_ids[None, ...]
+
+ audio_scales = model_kwargs.get("audio_scales")
+ if audio_scales is None:
+ audio_scales = [None] * batch_size
+
+ if self.decoder.config.audio_channels == 1:
+ output_values = self.audio_encoder.decode(
+ output_ids,
+ audio_scales=audio_scales,
+ ).audio_values
+ else:
+ codec_outputs_left = self.audio_encoder.decode(output_ids[:, :, ::2, :], audio_scales=audio_scales)
+ output_values_left = codec_outputs_left.audio_values
+
+ codec_outputs_right = self.audio_encoder.decode(output_ids[:, :, 1::2, :], audio_scales=audio_scales)
+ output_values_right = codec_outputs_right.audio_values
+
+ output_values = torch.cat([output_values_left, output_values_right], dim=1)
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_values
+ return outputs
+ else:
+ return output_values
+
+ def get_unconditional_inputs(self, num_samples=1):
+ """
+ Helper function to get null inputs for unconditional generation, enabling the model to be used without the
+ feature extractor or tokenizer.
+
+ Args:
+ num_samples (int, *optional*):
+ Number of audio samples to unconditionally generate.
+ max_new_tokens (int, *optional*):
+ Number of tokens to generate for each sample. More tokens means longer audio samples, at the expense of
+ longer inference (since more audio tokens need to be generated per sample).
+
+ Example:
+ ```python
+ >>> from transformers import MusicgenForConditionalGeneration
+
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+ >>> # get the unconditional (or 'null') inputs for the model
+ >>> unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
+ >>> audio_samples = model.generate(**unconditional_inputs, max_new_tokens=256)
+ ```"""
+ last_hidden_state = torch.zeros(
+ (num_samples, 1, self.config.text_encoder.hidden_size), device=self.device, dtype=self.dtype
+ )
+
+ attention_mask = torch.zeros((num_samples, 1), device=self.device, dtype=torch.long)
+
+ return MusicgenUnconditionalInput(
+ encoder_outputs=(last_hidden_state,),
+ attention_mask=attention_mask,
+ guidance_scale=1.0,
+ )
+
+
+__all__ = ["MusicgenForConditionalGeneration", "MusicgenForCausalLM", "MusicgenModel", "MusicgenPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py b/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..82671e4bf4cc25573528a9a849780fa7b3b8a5f8
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py
@@ -0,0 +1,144 @@
+# 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.
+"""
+Text/audio processor class for MusicGen
+"""
+
+from typing import List, Optional
+
+import numpy as np
+
+from ...processing_utils import ProcessorMixin
+from ...utils import to_numpy
+
+
+class MusicgenProcessor(ProcessorMixin):
+ r"""
+ Constructs a MusicGen processor which wraps an EnCodec feature extractor and a T5 tokenizer into a single processor
+ class.
+
+ [`MusicgenProcessor`] offers all the functionalities of [`EncodecFeatureExtractor`] and [`TTokenizer`]. See
+ [`~MusicgenProcessor.__call__`] and [`~MusicgenProcessor.decode`] for more information.
+
+ Args:
+ feature_extractor (`EncodecFeatureExtractor`):
+ An instance of [`EncodecFeatureExtractor`]. The feature extractor is a required input.
+ tokenizer (`T5Tokenizer`):
+ An instance of [`T5Tokenizer`]. The tokenizer is a required input.
+ """
+
+ feature_extractor_class = "EncodecFeatureExtractor"
+ tokenizer_class = ("T5Tokenizer", "T5TokenizerFast")
+
+ def __init__(self, feature_extractor, tokenizer):
+ super().__init__(feature_extractor, tokenizer)
+ self.current_processor = self.feature_extractor
+ self._in_target_context_manager = False
+
+ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
+ return self.tokenizer.get_decoder_prompt_ids(task=task, language=language, no_timestamps=no_timestamps)
+
+ def __call__(self, *args, **kwargs):
+ """
+ Forwards the `audio` argument to EncodecFeatureExtractor's [`~EncodecFeatureExtractor.__call__`] and the `text`
+ argument to [`~T5Tokenizer.__call__`]. 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(*args, **kwargs)
+
+ 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 text is not None:
+ inputs = self.tokenizer(text, **kwargs)
+
+ if audio is not None:
+ audio_inputs = self.feature_extractor(audio, *args, sampling_rate=sampling_rate, **kwargs)
+
+ if audio is None:
+ return inputs
+
+ elif text is None:
+ return audio_inputs
+
+ else:
+ inputs["input_values"] = audio_inputs["input_values"]
+ if "padding_mask" in audio_inputs:
+ inputs["padding_mask"] = audio_inputs["padding_mask"]
+ return inputs
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method is used to decode either batches of audio outputs from the MusicGen model, or batches of token ids
+ from the tokenizer. In the case of decoding token ids, this method forwards all its arguments to T5Tokenizer's
+ [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information.
+ """
+ audio_values = kwargs.pop("audio", None)
+ padding_mask = kwargs.pop("padding_mask", None)
+
+ if len(args) > 0:
+ audio_values = args[0]
+ args = args[1:]
+
+ if audio_values is not None:
+ return self._decode_audio(audio_values, padding_mask=padding_mask)
+ else:
+ return self.tokenizer.batch_decode(*args, **kwargs)
+
+ def decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to T5Tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to the
+ docstring of this method for more information.
+ """
+ return self.tokenizer.decode(*args, **kwargs)
+
+ def _decode_audio(self, audio_values, padding_mask: Optional = None) -> List[np.ndarray]:
+ """
+ This method strips any padding from the audio values to return a list of numpy audio arrays.
+ """
+ audio_values = to_numpy(audio_values)
+ bsz, channels, seq_len = audio_values.shape
+
+ if padding_mask is None:
+ return list(audio_values)
+
+ padding_mask = to_numpy(padding_mask)
+
+ # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
+ # token (so that the generated audio values are **not** treated as padded tokens)
+ difference = seq_len - padding_mask.shape[-1]
+ padding_value = 1 - self.feature_extractor.padding_value
+ padding_mask = np.pad(padding_mask, ((0, 0), (0, difference)), "constant", constant_values=padding_value)
+
+ audio_values = audio_values.tolist()
+ for i in range(bsz):
+ sliced_audio = np.asarray(audio_values[i])[
+ padding_mask[i][None, :] != self.feature_extractor.padding_value
+ ]
+ audio_values[i] = sliced_audio.reshape(channels, -1)
+
+ return audio_values
+
+
+__all__ = ["MusicgenProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/__init__.py b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..20c8507aaed7b3922dbb71c054c6a8074edc2b7a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__init__.py
@@ -0,0 +1,86 @@
+# 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,
+ is_torchaudio_available,
+)
+
+
+_import_structure = {
+ "configuration_musicgen_melody": [
+ "MusicgenMelodyConfig",
+ "MusicgenMelodyDecoderConfig",
+ ],
+}
+
+try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ pass
+else:
+ _import_structure["modeling_musicgen_melody"] = [
+ "MusicgenMelodyForConditionalGeneration",
+ "MusicgenMelodyForCausalLM",
+ "MusicgenMelodyModel",
+ "MusicgenMelodyPreTrainedModel",
+ ]
+
+try:
+ if not is_torchaudio_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ pass
+else:
+ _import_structure["feature_extraction_musicgen_melody"] = ["MusicgenMelodyFeatureExtractor"]
+ _import_structure["processing_musicgen_melody"] = ["MusicgenMelodyProcessor"]
+
+
+if TYPE_CHECKING:
+ from .configuration_musicgen_melody import (
+ MusicgenMelodyConfig,
+ MusicgenMelodyDecoderConfig,
+ )
+
+ try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ pass
+ else:
+ from .modeling_musicgen_melody import (
+ MusicgenMelodyForCausalLM,
+ MusicgenMelodyForConditionalGeneration,
+ MusicgenMelodyModel,
+ MusicgenMelodyPreTrainedModel,
+ )
+
+ try:
+ if not is_torchaudio_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ pass
+ else:
+ from .feature_extraction_musicgen_melody import MusicgenMelodyFeatureExtractor
+ from .processing_musicgen_melody import MusicgenMelodyProcessor
+
+
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..df26f87fcbf6812eaedffa355ecc138075c6037e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/configuration_musicgen_melody.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/configuration_musicgen_melody.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..38f98b207a9c57fda23a9ea487173ccca3b90e9b
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/configuration_musicgen_melody.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/feature_extraction_musicgen_melody.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/feature_extraction_musicgen_melody.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..854531e300284edb5db0216f4197459f3ea12aff
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/feature_extraction_musicgen_melody.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/processing_musicgen_melody.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/processing_musicgen_melody.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a13994c012792817ea126203cfbdee56aee58bf4
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/musicgen_melody/__pycache__/processing_musicgen_melody.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/configuration_musicgen_melody.py b/lib/python3.12/site-packages/transformers/models/musicgen_melody/configuration_musicgen_melody.py
new file mode 100644
index 0000000000000000000000000000000000000000..e65ad50021c3abf343b4da54396bad58aded0385
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen_melody/configuration_musicgen_melody.py
@@ -0,0 +1,258 @@
+# coding=utf-8
+# Copyright 2024 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.
+"""Musicgen Melody model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ..auto.configuration_auto import AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class MusicgenMelodyDecoderConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of an [`MusicgenMelodyDecoder`]. It is used to instantiate a
+ Musicgen Melody decoder according to the specified arguments, defining the model architecture. Instantiating a
+ configuration with the defaults will yield a similar configuration to that of the Musicgen Melody
+ [facebook/musicgen-melody](https://huggingface.co/facebook/musicgen-melody) 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 2048):
+ Vocabulary size of the MusicgenMelodyDecoder model. Defines the number of different tokens that can be
+ represented by the `inputs_ids` passed when calling [`MusicgenMelodyDecoder`].
+ 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).
+ num_hidden_layers (`int`, *optional*, defaults to 24):
+ Number of decoder layers.
+ ffn_dim (`int`, *optional*, defaults to 4096):
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer block.
+ num_attention_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads for each attention layer in the Transformer block.
+ 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.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether the model should return the last key/values attentions (not used by all models)
+ activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the decoder and pooler. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ hidden_size (`int`, *optional*, defaults to 1024):
+ Dimensionality of the layers and the pooler layer.
+ dropout (`float`, *optional*, defaults to 0.1):
+ The dropout probability for all fully connected layers in the embeddings, text_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.
+ initializer_factor (`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 `False`):
+ Scale embeddings by diving by sqrt(hidden_size).
+ num_codebooks (`int`, *optional*, defaults to 4):
+ The number of parallel codebooks forwarded to the model.
+ audio_channels (`int`, *optional*, defaults to 1):
+ Number of audio channels used by the model (either mono or stereo). Stereo models generate a separate
+ audio stream for the left/right output channels. Mono models generate a single audio stream output.
+ pad_token_id (`int`, *optional*, defaults to 2048): The id of the *padding* token.
+ bos_token_id (`int`, *optional*, defaults to 2048): The id of the *beginning-of-sequence* token.
+ eos_token_id (`int`, *optional*): The id of the *end-of-sequence* token.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie word embeddings with the text encoder.
+ """
+
+ model_type = "musicgen_melody_decoder"
+ base_config_key = "decoder_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=2048,
+ max_position_embeddings=2048,
+ num_hidden_layers=24,
+ ffn_dim=4096,
+ num_attention_heads=16,
+ layerdrop=0.0,
+ use_cache=True,
+ activation_function="gelu",
+ hidden_size=1024,
+ dropout=0.1,
+ attention_dropout=0.0,
+ activation_dropout=0.0,
+ initializer_factor=0.02,
+ scale_embedding=False,
+ num_codebooks=4,
+ audio_channels=1,
+ pad_token_id=2048,
+ bos_token_id=2048,
+ eos_token_id=None,
+ tie_word_embeddings=False,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.ffn_dim = ffn_dim
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.dropout = dropout
+ self.attention_dropout = attention_dropout
+ self.activation_dropout = activation_dropout
+ self.activation_function = activation_function
+ self.initializer_factor = initializer_factor
+ self.layerdrop = layerdrop
+ self.use_cache = use_cache
+ self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
+ self.num_codebooks = num_codebooks
+
+ if audio_channels not in [1, 2]:
+ raise ValueError(f"Expected 1 (mono) or 2 (stereo) audio channels, got {audio_channels} channels.")
+ self.audio_channels = audio_channels
+
+ 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,
+ )
+
+
+class MusicgenMelodyConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`MusicgenMelodyModel`]. It is used to instantiate a
+ Musicgen Melody model according to the specified arguments, defining the text encoder, audio encoder and Musicgen Melody decoder
+ configs. Instantiating a configuration with the defaults will yield a similar configuration to that of the Musicgen Melody
+ [facebook/musicgen-melody](https://huggingface.co/facebook/musicgen-melody) 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_chroma (`int`, *optional*, defaults to 12): Number of chroma bins to use.
+ chroma_length (`int`, *optional*, defaults to 235):
+ Maximum chroma duration if audio is used to condition the model. Corresponds to the maximum duration used during training.
+ kwargs (*optional*):
+ Dictionary of keyword arguments. Notably:
+
+ - **text_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+ defines the text encoder config.
+ - **audio_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+ defines the audio encoder config.
+ - **decoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that defines
+ the decoder config.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... MusicgenMelodyConfig,
+ ... MusicgenMelodyDecoderConfig,
+ ... T5Config,
+ ... EncodecConfig,
+ ... MusicgenMelodyForConditionalGeneration,
+ ... )
+
+ >>> # Initializing text encoder, audio encoder, and decoder model configurations
+ >>> text_encoder_config = T5Config()
+ >>> audio_encoder_config = EncodecConfig()
+ >>> decoder_config = MusicgenMelodyDecoderConfig()
+
+ >>> configuration = MusicgenMelodyConfig.from_sub_models_config(
+ ... text_encoder_config, audio_encoder_config, decoder_config
+ ... )
+
+ >>> # Initializing a MusicgenMelodyForConditionalGeneration (with random weights) from the facebook/musicgen-melody style configuration
+ >>> model = MusicgenMelodyForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ >>> config_text_encoder = model.config.text_encoder
+ >>> config_audio_encoder = model.config.audio_encoder
+ >>> config_decoder = model.config.decoder
+
+ >>> # Saving the model, including its configuration
+ >>> model.save_pretrained("musicgen_melody-model")
+
+ >>> # loading model and config from pretrained folder
+ >>> musicgen_melody_config = MusicgenMelodyConfig.from_pretrained("musicgen_melody-model")
+ >>> model = MusicgenMelodyForConditionalGeneration.from_pretrained("musicgen_melody-model", config=musicgen_melody_config)
+ ```"""
+
+ model_type = "musicgen_melody"
+ sub_configs = {
+ "text_encoder": AutoConfig,
+ "audio_encoder": AutoConfig,
+ "decoder": MusicgenMelodyDecoderConfig,
+ }
+ is_composition = True
+
+ def __init__(
+ self,
+ num_chroma=12,
+ chroma_length=235,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
+ raise ValueError("Config has to be initialized with text_encoder, audio_encoder and decoder config")
+
+ text_encoder_config = kwargs.pop("text_encoder")
+ text_encoder_model_type = text_encoder_config.pop("model_type")
+
+ audio_encoder_config = kwargs.pop("audio_encoder")
+ audio_encoder_model_type = audio_encoder_config.pop("model_type")
+
+ decoder_config = kwargs.pop("decoder")
+
+ self.text_encoder = AutoConfig.for_model(text_encoder_model_type, **text_encoder_config)
+ self.audio_encoder = AutoConfig.for_model(audio_encoder_model_type, **audio_encoder_config)
+ self.decoder = MusicgenMelodyDecoderConfig(**decoder_config)
+ self.is_encoder_decoder = False
+
+ self.num_chroma = num_chroma
+ self.chroma_length = chroma_length
+
+ @classmethod
+ def from_sub_models_config(
+ cls,
+ text_encoder_config: PretrainedConfig,
+ audio_encoder_config: PretrainedConfig,
+ decoder_config: MusicgenMelodyDecoderConfig,
+ **kwargs,
+ ):
+ r"""
+ Instantiate a [`MusicgenMelodyConfig`] (or a derived class) from text encoder, audio encoder and decoder
+ configurations.
+
+ Returns:
+ [`MusicgenMelodyConfig`]: An instance of a configuration object
+ """
+
+ return cls(
+ text_encoder=text_encoder_config.to_dict(),
+ audio_encoder=audio_encoder_config.to_dict(),
+ decoder=decoder_config.to_dict(),
+ **kwargs,
+ )
+
+ @property
+ # This is a property because you might want to change the codec model on the fly
+ def sampling_rate(self):
+ return self.audio_encoder.sampling_rate
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/feature_extraction_musicgen_melody.py b/lib/python3.12/site-packages/transformers/models/musicgen_melody/feature_extraction_musicgen_melody.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec490b7d9048987d8b9f269aac19f6f9d8565ee3
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen_melody/feature_extraction_musicgen_melody.py
@@ -0,0 +1,331 @@
+# coding=utf-8
+# Copyright 2024 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.
+"""
+Feature extractor class for Musicgen Melody
+"""
+
+import copy
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from ...audio_utils import chroma_filter_bank
+from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
+from ...feature_extraction_utils import BatchFeature
+from ...utils import TensorType, is_torch_available, is_torchaudio_available, logging
+
+
+if is_torch_available():
+ import torch
+
+if is_torchaudio_available():
+ import torchaudio
+
+logger = logging.get_logger(__name__)
+
+
+class MusicgenMelodyFeatureExtractor(SequenceFeatureExtractor):
+ r"""
+ Constructs a MusicgenMelody 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 class extracts chroma features from audio processed by [Demucs](https://github.com/adefossez/demucs/tree/main) or
+ directly from raw audio waveform.
+
+ Args:
+ feature_size (`int`, *optional*, defaults to 12):
+ The feature dimension of the extracted features.
+ sampling_rate (`int`, *optional*, defaults to 32000):
+ The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
+ hop_length (`int`, *optional*, defaults to 4096):
+ Length of the overlaping windows for the STFT used to obtain the Mel Frequency coefficients.
+ chunk_length (`int`, *optional*, defaults to 30):
+ The maximum number of chunks of `sampling_rate` samples used to trim and pad longer or shorter audio
+ sequences.
+ n_fft (`int`, *optional*, defaults to 16384):
+ Size of the Fourier transform.
+ num_chroma (`int`, *optional*, defaults to 12):
+ Number of chroma bins to use.
+ padding_value (`float`, *optional*, defaults to 0.0):
+ Padding value used to pad the audio.
+ return_attention_mask (`bool`, *optional*, defaults to `False`):
+ Whether to return the attention mask. Can be overwritten when calling the feature extractor.
+
+ [What are attention masks?](../glossary#attention-mask)
+
+
+
+ For Whisper models, `attention_mask` should always be passed for batched inference, to avoid subtle
+ bugs.
+
+
+ stem_indices (`List[int]`, *optional*, defaults to `[3, 2]`):
+ Stem channels to extract if demucs outputs are passed.
+ """
+
+ model_input_names = ["input_features"]
+
+ def __init__(
+ self,
+ feature_size=12,
+ sampling_rate=32000,
+ hop_length=4096,
+ chunk_length=30,
+ n_fft=16384,
+ num_chroma=12,
+ padding_value=0.0,
+ return_attention_mask=False, # pad inputs to max length with silence token (zero) and no attention mask
+ stem_indices=[3, 2],
+ **kwargs,
+ ):
+ super().__init__(
+ feature_size=feature_size,
+ sampling_rate=sampling_rate,
+ padding_value=padding_value,
+ return_attention_mask=return_attention_mask,
+ **kwargs,
+ )
+ self.n_fft = n_fft
+ self.hop_length = hop_length
+ self.chunk_length = chunk_length
+ self.n_samples = chunk_length * sampling_rate
+ self.sampling_rate = sampling_rate
+ self.chroma_filters = torch.from_numpy(
+ chroma_filter_bank(sampling_rate=sampling_rate, num_frequency_bins=n_fft, tuning=0, num_chroma=num_chroma)
+ ).float()
+ self.spectrogram = torchaudio.transforms.Spectrogram(
+ n_fft=n_fft, win_length=n_fft, hop_length=hop_length, power=2, center=True, pad=0, normalized=True
+ )
+ self.stem_indices = stem_indices
+
+ def _torch_extract_fbank_features(self, waveform: torch.Tensor) -> torch.Tensor:
+ """
+ Compute the chroma spectrogram of the provided audio using the torchaudio spectrogram implementation and the librosa chroma features.
+ """
+
+ # if wav length is not long enough, pad it
+ wav_length = waveform.shape[-1]
+ if wav_length < self.n_fft:
+ pad = self.n_fft - wav_length
+ rest = 0 if pad % 2 == 0 else 1
+ waveform = torch.nn.functional.pad(waveform, (pad // 2, pad // 2 + rest), "constant", 0)
+
+ # squeeze alongside channel dimension
+ spec = self.spectrogram(waveform).squeeze(1)
+
+ # sum along the frequency dimension
+ raw_chroma = torch.einsum("cf, ...ft->...ct", self.chroma_filters, spec)
+
+ # normalise with max value
+ norm_chroma = torch.nn.functional.normalize(raw_chroma, p=float("inf"), dim=-2, eps=1e-6)
+
+ # transpose time and chroma dimension -> (batch, time, chroma)
+ norm_chroma = norm_chroma.transpose(1, 2)
+
+ # replace max value alongside chroma dimension with 1 and replace the rest with 0
+ idx = norm_chroma.argmax(-1, keepdim=True)
+ norm_chroma[:] = 0
+ norm_chroma.scatter_(dim=-1, index=idx, value=1)
+
+ return norm_chroma
+
+ def _extract_stem_indices(self, audio, sampling_rate=None):
+ """
+ Extracts stems from the output of the [Demucs](https://github.com/adefossez/demucs/tree/main) audio separation model,
+ then converts to mono-channel and resample to the feature extractor sampling rate.
+
+ Args:
+ audio (`torch.Tensor` of shape `(batch_size, num_stems, channel_size, audio_length)`):
+ The output of the Demucs model to be processed.
+ sampling_rate (`int`, *optional*):
+ Demucs sampling rate. If not specified, defaults to `44000`.
+ """
+ sampling_rate = 44000 if sampling_rate is None else sampling_rate
+
+ # extract "vocals" and "others" sources from audio encoder (demucs) output
+ # [batch_size, num_stems, channel_size, audio_length]
+ wav = audio[:, torch.tensor(self.stem_indices)]
+
+ # merge extracted stems to single waveform
+ wav = wav.sum(1)
+
+ # convert to mono-channel waveform
+ wav = wav.mean(dim=1, keepdim=True)
+
+ # resample to model sampling rate
+ # not equivalent to julius.resample
+ if sampling_rate != self.sampling_rate:
+ wav = torchaudio.functional.resample(
+ wav, sampling_rate, self.sampling_rate, rolloff=0.945, lowpass_filter_width=24
+ )
+
+ # [batch_size, 1, audio_length] -> [batch_size, audio_length]
+ wav = wav.squeeze(1)
+
+ return wav
+
+ def __call__(
+ self,
+ audio: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
+ truncation: bool = True,
+ pad_to_multiple_of: Optional[int] = None,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ return_attention_mask: Optional[bool] = None,
+ padding: Optional[str] = True,
+ max_length: Optional[int] = None,
+ sampling_rate: Optional[int] = None,
+ **kwargs,
+ ) -> BatchFeature:
+ """
+ Main method to featurize and prepare for the model one or several sequence(s).
+
+ Args:
+ audio (`torch.Tensor`, `np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[torch.Tensor]`, `List[List[float]]`):
+ The sequence or batch of sequences to be padded. Each sequence can be a torch tensor, a numpy array, a list of float
+ values, a list of numpy arrays, a list of torch tensors, or a list of list of float values.
+ If `audio` is the output of Demucs, it has to be a torch tensor of shape `(batch_size, num_stems, channel_size, audio_length)`.
+ Otherwise, it must be mono or stereo channel audio.
+ truncation (`bool`, *optional*, default to `True`):
+ Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+ pad_to_multiple_of (`int`, *optional*, defaults to None):
+ 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_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 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)
+
+
+ For Musicgen Melody models, audio `attention_mask` is not necessary.
+
+
+ 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).
+ max_length (`int`, *optional*):
+ Maximum length of the returned list and optionally padding length (see above).
+ 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.
+ Note that if `audio` is the output of Demucs, `sampling_rate` must be the sampling rate at which Demucs operates.
+ """
+
+ if sampling_rate is None:
+ logger.warning_once(
+ f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
+ "Failing to do so can result in silent errors that might be hard to debug."
+ )
+
+ if isinstance(audio, torch.Tensor) and len(audio.shape) == 4:
+ logger.warning_once(
+ "`audio` is a 4-dimensional torch tensor and has thus been recognized as the output of `Demucs`. "
+ "If this is not the case, make sure to read Musicgen Melody docstrings and "
+ "to correct `audio` to get the right behaviour."
+ "Link to the docstrings: https://huggingface.co/docs/transformers/main/en/model_doc/musicgen_melody"
+ )
+ audio = self._extract_stem_indices(audio, sampling_rate=sampling_rate)
+ elif sampling_rate is not None and sampling_rate != self.sampling_rate:
+ audio = torchaudio.functional.resample(
+ audio, sampling_rate, self.sampling_rate, rolloff=0.945, lowpass_filter_width=24
+ )
+
+ is_batched = isinstance(audio, (np.ndarray, torch.Tensor)) and len(audio.shape) > 1
+ is_batched = is_batched or (
+ isinstance(audio, (list, tuple)) and (isinstance(audio[0], (torch.Tensor, np.ndarray, tuple, list)))
+ )
+
+ if is_batched and not isinstance(audio[0], torch.Tensor):
+ audio = [torch.tensor(speech, dtype=torch.float32).unsqueeze(-1) for speech in audio]
+ elif is_batched:
+ audio = [speech.unsqueeze(-1) for speech in audio]
+ elif not is_batched and not isinstance(audio, torch.Tensor):
+ audio = torch.tensor(audio, dtype=torch.float32).unsqueeze(-1)
+
+ if isinstance(audio[0], torch.Tensor) and audio[0].dtype is torch.float64:
+ audio = [speech.to(torch.float32) for speech in audio]
+
+ # always return batch
+ if not is_batched:
+ audio = [audio]
+
+ if len(audio[0].shape) == 3:
+ logger.warning_once(
+ "`audio` has been detected as a batch of stereo signals. Will be convert to mono signals. "
+ "If this is an undesired behaviour, make sure to read Musicgen Melody docstrings and "
+ "to correct `audio` to get the right behaviour."
+ "Link to the docstrings: https://huggingface.co/docs/transformers/main/en/model_doc/musicgen_melody"
+ )
+ # convert to mono-channel waveform
+ audio = [stereo.mean(dim=0) for stereo in audio]
+
+ batched_speech = BatchFeature({"input_features": audio})
+
+ padded_inputs = self.pad(
+ batched_speech,
+ padding=padding,
+ max_length=max_length if max_length else self.n_samples,
+ truncation=truncation,
+ pad_to_multiple_of=pad_to_multiple_of,
+ return_attention_mask=return_attention_mask,
+ return_tensors="pt",
+ )
+
+ input_features = self._torch_extract_fbank_features(padded_inputs["input_features"].squeeze(-1))
+
+ padded_inputs["input_features"] = input_features
+
+ if return_attention_mask:
+ # rescale from raw audio length to spectrogram length
+ padded_inputs["attention_mask"] = padded_inputs["attention_mask"][:, :: self.hop_length]
+
+ if return_tensors is not None:
+ padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
+
+ return padded_inputs
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ output = copy.deepcopy(self.__dict__)
+ output["feature_extractor_type"] = self.__class__.__name__
+ if "mel_filters" in output:
+ del output["mel_filters"]
+ if "window" in output:
+ del output["window"]
+ if "chroma_filters" in output:
+ del output["chroma_filters"]
+ if "spectrogram" in output:
+ del output["spectrogram"]
+ return output
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/modeling_musicgen_melody.py b/lib/python3.12/site-packages/transformers/models/musicgen_melody/modeling_musicgen_melody.py
new file mode 100644
index 0000000000000000000000000000000000000000..70b914313c89bd3aa4122bbacd191bfa5502aeb1
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen_melody/modeling_musicgen_melody.py
@@ -0,0 +1,2565 @@
+# coding=utf-8
+# Copyright 2024 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 Musicgen Melody model."""
+
+import copy
+import inspect
+import math
+import random
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from ...activations import ACT2FN
+from ...generation import (
+ ClassifierFreeGuidanceLogitsProcessor,
+ GenerationConfig,
+ GenerationMixin,
+ GenerationMode,
+ LogitsProcessorList,
+ StoppingCriteriaList,
+)
+from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_outputs import (
+ BaseModelOutputWithPast,
+ ModelOutput,
+)
+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, AutoModelForTextEncoding
+from .configuration_musicgen_melody import MusicgenMelodyConfig, MusicgenMelodyDecoderConfig
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+if TYPE_CHECKING:
+ from ...generation.streamers import BaseStreamer
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MusicgenMelodyConfig"
+_CHECKPOINT_FOR_DOC = "facebook/musicgen-melody"
+
+
+@dataclass
+class MusicgenMelodyOutputWithPast(ModelOutput):
+ """
+ Base class for Musicgen Melody autoregressive outputs.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token 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).
+ 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)`)
+
+ 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.
+ 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.
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of conditional hidden-states representing the concatenation of the projeted text encoder output and the projeted audio encoder output.
+ Used as a conditional signal.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[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
+ encoder_hidden_states: Optional[torch.FloatTensor] = None
+
+
+# Copied from transformers.models.musicgen.modeling_musicgen.shift_tokens_right
+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.
+ """
+ # transpose to get (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.transpose(1, 2)
+ 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
+
+
+# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenSinusoidalPositionalEmbedding with Musicgen->MusicgenMelody
+class MusicgenMelodySinusoidalPositionalEmbedding(nn.Module):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, num_positions: int, embedding_dim: int):
+ super().__init__()
+ self.embedding_dim = embedding_dim
+ self.make_weights(num_positions, embedding_dim)
+
+ def make_weights(self, num_embeddings: int, embedding_dim: int):
+ emb_weights = self.get_embedding(num_embeddings, embedding_dim)
+ if hasattr(self, "weights"):
+ # in forward put the weights on the correct dtype and device of the param
+ emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
+
+ self.weights = nn.Parameter(emb_weights)
+ self.weights.requires_grad = False
+ self.weights.detach_()
+
+ @staticmethod
+ def get_embedding(num_embeddings: int, embedding_dim: int):
+ """
+ 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.cos(emb), torch.sin(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)
+ return emb.to(torch.get_default_dtype())
+
+ @torch.no_grad()
+ # Ignore copy
+ def forward(self, inputs_embeds: torch.Tensor, past_key_values_length: int = 0):
+ bsz, seq_len, _ = inputs_embeds.size()
+ # Create the position ids from the input token ids.
+ position_ids = (torch.arange(seq_len) + past_key_values_length).to(inputs_embeds.device)
+ # expand embeddings if needed
+ if seq_len > self.weights.size(0):
+ self.make_weights(seq_len + self.offset, self.embedding_dim)
+ return self.weights.index_select(0, position_ids.view(-1)).detach()
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->MusicgenMelody
+class MusicgenMelodyAttention(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[MusicgenMelodyConfig] = 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.bart.modeling_bart.BartFlashAttention2 with Bart->MusicgenMelody
+class MusicgenMelodyFlashAttention2(MusicgenMelodyAttention):
+ """
+ MusicgenMelody flash attention module. This module inherits from `MusicgenMelodyAttention` 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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ def _reshape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+
+ 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]]]:
+ # MusicgenMelodyFlashAttention2 attention does not support output_attentions
+ if output_attentions:
+ raise ValueError("MusicgenMelodyFlashAttention2 attention does not support output_attentions")
+
+ # 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, q_len, _ = hidden_states.size()
+
+ # get query proj
+ query_states = self._reshape(self.q_proj(hidden_states), -1, bsz)
+ # 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].transpose(1, 2)
+ value_states = past_key_value[1].transpose(1, 2)
+ elif is_cross_attention:
+ # cross_attentions
+ key_states = self._reshape(self.k_proj(key_value_states), -1, bsz)
+ value_states = self._reshape(self.v_proj(key_value_states), -1, bsz)
+ elif past_key_value is not None:
+ # reuse k, v, self_attention
+ key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)
+ key_states = torch.cat([past_key_value[0].transpose(1, 2), key_states], dim=1)
+ value_states = torch.cat([past_key_value[1].transpose(1, 2), value_states], dim=1)
+ else:
+ # self_attention
+ key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._reshape(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.transpose(1, 2), value_states.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]
+
+ # 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 = 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 = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=self.dropout if self.training else 0.0,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1)
+ 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.bart.modeling_bart.BartSdpaAttention with Bart->MusicgenMelody
+class MusicgenMelodySdpaAttention(MusicgenMelodyAttention):
+ 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 output_attentions or layer_head_mask is not None:
+ # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.
+ logger.warning_once(
+ "MusicgenMelodyModel is using MusicgenMelodySdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_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.'
+ )
+ return super().forward(
+ hidden_states,
+ key_value_states=key_value_states,
+ past_key_value=past_key_value,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ output_attentions=output_attentions,
+ )
+
+ # 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)
+ # 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)
+
+ query_states = self._shape(query_states, tgt_len, bsz)
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case tgt_len == 1.
+ is_causal = True if self.is_causal and attention_mask is None and tgt_len > 1 else False
+
+ # NOTE: SDPA with memory-efficient backend is currently (torch==2.1.2) bugged when using non-contiguous inputs and a custom attn_mask,
+ # but we are fine here as `_shape` do call `.contiguous()`. Reference: https://github.com/pytorch/pytorch/issues/112577
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attn_mask=attention_mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ is_causal=is_causal,
+ )
+
+ 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.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, None, past_key_value
+
+
+MUSICGEN_MELODY_ATTENTION_CLASSES = {
+ "eager": MusicgenMelodyAttention,
+ "sdpa": MusicgenMelodySdpaAttention,
+ "flash_attention_2": MusicgenMelodyFlashAttention2,
+}
+
+
+class MusicgenMelodyDecoderLayer(nn.Module):
+ def __init__(self, config: MusicgenMelodyDecoderConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+
+ self.self_attn = MUSICGEN_MELODY_ATTENTION_CLASSES[config._attn_implementation](
+ embed_dim=self.embed_dim,
+ num_heads=config.num_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ is_causal=True,
+ config=config,
+ )
+ 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.fc1 = nn.Linear(self.embed_dim, config.ffn_dim, bias=False)
+ self.fc2 = nn.Linear(config.ffn_dim, self.embed_dim, bias=False)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ layer_head_mask: Optional[torch.Tensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = True,
+ ) -> torch.Tensor:
+ """
+ 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.
+ layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(attention_heads,)`.
+ past_key_value (`Tuple(torch.FloatTensor)`): 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.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # 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
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_value=self_attn_past_key_value,
+ attention_mask=attention_mask,
+ layer_head_mask=layer_head_mask,
+ output_attentions=output_attentions,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(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
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights,)
+
+ if use_cache:
+ outputs += (present_key_value,)
+
+ return outputs
+
+
+# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenPreTrainedModel with Musicgen->MusicgenMelody
+class MusicgenMelodyPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = MusicgenMelodyDecoderConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MusicgenMelodyDecoderLayer", "MusicgenMelodyAttention"]
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_factor
+ 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_()
+ 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, MusicgenMelodySinusoidalPositionalEmbedding):
+ weights = module.get_embedding(*module.weights.shape)
+ weights = nn.Parameter(weights, requires_grad=False)
+ weights.detach_()
+ module.weights = weights
+
+
+MUSICGEN_MELODY_START_DOCSTRING = r"""
+
+ The Musicgen Melody model was proposed in [Simple and Controllable Music Generation](https://arxiv.org/abs/2306.05284) by
+ Jade Copet, Felix Kreuk, Itai Gat, Tal Remez, David Kant, Gabriel Synnaeve, Yossi Adi, Alexandre Défossez. It is a
+ decoder-only transformer trained on the task of conditional music generation.
+
+ 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 ([`MusicgenMelodyConfig`]): 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.
+"""
+
+MUSICGEN_MELODY_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)
+ input_features (`torch.FloatTensor` of shape `(batch_size, audio_sequence_length, num_chroma)`):
+ Input audio features.
+ This should be returned by the [`MusicgenMelodyFeatureExtractor`] class that you can also
+ retrieve from [`AutoFeatureExtractor`]. See [`MusicgenMelodyFeatureExtractor.__call__`] for details.
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+
+
+ The `decoder_input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `decoder_input_ids`.
+
+
+
+ decoder_attention_mask (`torch.LongTensor` 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.
+ 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, encoder_sequence_length + sequence_length, embed_size_per_head)`).
+
+ 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)`.
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of conditional hidden-states representing the concatenation of the projeted text encoder output and the projeted audio encoder output.
+ Used as a conditional signal and will thus be concatenated to the projeted `decoder_input_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.
+ 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`.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *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]`
+ 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.
+"""
+
+MUSICGEN_MELODY_DECODER_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+
+
+
+ The `input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `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)
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states representing the concatenation of the text encoder output and the processed audio encoder output.
+ Used as a conditional signal and will thus be concatenated to the projeted `decoder_input_ids`.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+ Mask to avoid performing attention on conditional hidden states. 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.Tensor` of shape `(decoder_layers, decoder_attention_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**.
+ 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)`)`.
+
+ 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)`.
+ 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.
+"""
+
+
+# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenDecoder with MUSICGEN->MUSICGEN_MELODY,Musicgen->MusicgenMelody
+class MusicgenMelodyDecoder(MusicgenMelodyPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MusicgenMelodyDecoderLayer`]
+ """
+
+ def __init__(self, config: MusicgenMelodyDecoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.layerdrop
+ self.max_target_positions = config.max_position_embeddings
+ self.d_model = config.hidden_size
+ self.num_codebooks = config.num_codebooks
+ self.embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
+
+ embed_dim = config.vocab_size + 1
+ self.embed_tokens = nn.ModuleList(
+ [nn.Embedding(embed_dim, config.hidden_size) for _ in range(config.num_codebooks)]
+ )
+
+ self.embed_positions = MusicgenMelodySinusoidalPositionalEmbedding(
+ config.max_position_embeddings,
+ config.hidden_size,
+ )
+
+ self.layers = nn.ModuleList([MusicgenMelodyDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.layer_norm = nn.LayerNorm(config.hidden_size)
+ self.attn_implementation = config._attn_implementation
+
+ 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(MUSICGEN_MELODY_DECODER_INPUTS_DOCSTRING)
+ # Ignore copy
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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:
+ # (bsz * codebooks, seq_len) -> (bsz, codebooks, seq_len)
+ input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input.shape
+ input_shape = (bsz, seq_len)
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ input = inputs_embeds[:, :, -1:]
+ else:
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+ # 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 inputs_embeds is None:
+ inputs_embeds = sum([self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks)])
+
+ if encoder_hidden_states is not None:
+ # take care of attention masks
+ if encoder_attention_mask is not None and attention_mask is None:
+ attention_mask = torch.ones(inputs_embeds.shape[:2], device=inputs_embeds.device)
+
+ if attention_mask is not None:
+ if encoder_attention_mask is None:
+ encoder_attention_mask = torch.ones(encoder_hidden_states.shape[:2], device=attention_mask.device)
+ attention_mask = torch.cat([encoder_attention_mask, attention_mask], dim=1)
+
+ # fuse encoder_hidden_states and inputs_embeds
+ inputs_embeds = torch.cat([encoder_hidden_states, inputs_embeds], dim=1)
+
+ input_shape = inputs_embeds.size()[:-1]
+
+ if self.attn_implementation == "flash_attention_2":
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ elif self.attn_implementation == "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.
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+ attention_mask,
+ input_shape,
+ inputs_embeds,
+ past_key_values_length,
+ )
+ else:
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask, input_shape, inputs_embeds, past_key_values_length
+ )
+
+ # embed positions
+ positions = self.embed_positions(inputs_embeds, past_key_values_length)
+
+ hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ 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
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+ next_decoder_cache = () if use_cache 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, 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:
+ all_hidden_states += (hidden_states,)
+ dropout_probability = random.uniform(0, 1)
+ if self.training and (dropout_probability < self.layerdrop):
+ continue
+
+ 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.forward,
+ hidden_states,
+ attention_mask,
+ head_mask[idx] if head_mask is not None else None,
+ None,
+ output_attentions,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ layer_head_mask=(head_mask[idx] if head_mask is not None else None),
+ 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_attentions += (layer_outputs[1],)
+
+ hidden_states = self.layer_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_attentions] if v is not None)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_attentions,
+ )
+
+
+@add_start_docstrings(
+ "The bare MusicgenMelody decoder model outputting raw hidden-states without any specific head on top.",
+ MUSICGEN_MELODY_START_DOCSTRING,
+)
+# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenModel with MUSICGEN->MUSICGEN_MELODY,Musicgen->MusicgenMelody
+class MusicgenMelodyModel(MusicgenMelodyPreTrainedModel):
+ def __init__(self, config: MusicgenMelodyDecoderConfig):
+ super().__init__(config)
+ self.decoder = MusicgenMelodyDecoder(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.decoder.embed_tokens = value
+
+ def get_decoder(self):
+ return self.decoder
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_MELODY_DECODER_INPUTS_DOCSTRING)
+ # Ignore copy
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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
+
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
+ decoder_outputs = self.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ head_mask=head_mask,
+ 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,
+ )
+
+ if not return_dict:
+ return decoder_outputs
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ hidden_states=decoder_outputs.hidden_states,
+ attentions=decoder_outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ "The Musicgen Melody decoder model with a language modelling head on top.",
+ MUSICGEN_MELODY_START_DOCSTRING,
+)
+# Copied from transformers.models.musicgen.modeling_musicgen.MusicgenForCausalLM with MUSICGEN->MUSICGEN_MELODY,Musicgen->MusicgenMelody,MusicGen->Musicgen Melody
+class MusicgenMelodyForCausalLM(MusicgenMelodyPreTrainedModel, GenerationMixin):
+ def __init__(self, config: MusicgenMelodyDecoderConfig):
+ super().__init__(config)
+
+ self.model = MusicgenMelodyModel(config)
+
+ self.num_codebooks = config.num_codebooks
+ self.lm_heads = nn.ModuleList(
+ [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_codebooks)]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_heads
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_heads = new_embeddings
+
+ def set_decoder(self, decoder):
+ self.model.decoder = decoder
+
+ def get_decoder(self):
+ return self.model.decoder
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_MELODY_DECODER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=MusicgenMelodyOutputWithPast, config_class=_CONFIG_FOR_DOC)
+ # Ignore copy
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.LongTensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Tuple[Tuple[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,
+ labels: Optional[torch.LongTensor] = None,
+ ) -> Union[Tuple, MusicgenMelodyOutputWithPast]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *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:
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if (labels is not None) and (input_ids is None and inputs_embeds is None):
+ input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.bos_token_id)
+
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ head_mask=head_mask,
+ 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]
+
+ lm_logits = torch.stack([head(hidden_states) for head in self.lm_heads], dim=1)
+
+ loss = None
+ if labels is not None:
+ # since encoder hidden states have been concatenated to the decoder hidden states,
+ # we take the last timestamps corresponding to labels
+ logits = lm_logits[:, :, -labels.shape[1] :]
+
+ loss_fct = CrossEntropyLoss()
+ loss = torch.zeros([], device=self.device)
+
+ # per codebook cross-entropy
+ # ref: https://github.com/facebookresearch/audiocraft/blob/69fea8b290ad1b4b40d28f92d1dfc0ab01dbab85/audiocraft/solvers/musicgen.py#L242-L243
+ # -100 labels are ignored
+ labels = labels.masked_fill(labels == self.config.pad_token_id, -100)
+
+ # per codebook cross-entropy
+ for codebook in range(self.config.num_codebooks):
+ codebook_logits = logits[:, codebook].contiguous().view(-1, logits.shape[-1])
+ codebook_labels = labels[..., codebook].contiguous().view(-1)
+ loss += loss_fct(codebook_logits, codebook_labels)
+
+ loss = loss / self.config.num_codebooks
+
+ # (bsz, num_codebooks, seq_len, vocab_size) -> (bsz * num_codebooks, seq_len, vocab_size)
+ lm_logits = lm_logits.reshape(-1, *lm_logits.shape[2:])
+
+ if not return_dict:
+ output = (lm_logits,) + outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return MusicgenMelodyOutputWithPast(
+ loss=loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ # Ignore copy
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ head_mask=None,
+ past_key_values=None,
+ use_cache=True,
+ delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if delay_pattern_mask is None:
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ input_ids = self.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ input_ids = input_ids.repeat((2, 1))
+ if attention_mask is not None:
+ attention_mask = attention_mask.repeat((2, 1))
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states = torch.concatenate(
+ [encoder_hidden_states, torch.zeros_like(encoder_hidden_states)], dim=0
+ )
+
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = torch.concatenate(
+ encoder_attention_mask, torch.zeros_like(encoder_attention_mask), dim=0
+ )
+
+ if past_key_values is not None:
+ input_ids = input_ids[:, -1:]
+
+ # we only want to use conditional signal in the 1st generation step but keeping the attention mask
+ encoder_hidden_states = None
+
+ return {
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "encoder_hidden_states": encoder_hidden_states,
+ "encoder_attention_mask": encoder_attention_mask,
+ "head_mask": head_mask,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ }
+
+ def build_delay_pattern_mask(
+ self, input_ids: torch.LongTensor, pad_token_id: int, max_length: Optional[int] = None
+ ):
+ """Build a delayed pattern mask to the input_ids. Each codebook is offset by the previous codebook by
+ one, giving a delayed pattern mask at the start of sequence and end of sequence. Take the example where there
+ are 4 codebooks and a max sequence length of 8, we have the delayed pattern mask of shape `(codebooks,
+ seq_len)`:
+ - [P, -1, -1, -1, -1, P, P, P]
+ - [P, P, -1, -1, -1, -1, P, P]
+ - [P, P, P, -1, -1, -1, -1, P]
+ - [P, P, P, P, -1, -1, -1, -1]
+ where P is the special padding token id and -1 indicates that the token is valid for prediction. If we include
+ a prompt (decoder input ids), the -1 positions indicate where new tokens should be predicted. Otherwise, the
+ mask is set to the value in the prompt:
+ - [P, a, b, -1, -1, P, P, P]
+ - [P, P, c, d, -1, -1, P, P]
+ - [P, P, P, e, f, -1, -1, P]
+ - [P, P, P, P, g, h, -1, -1]
+ where a-h indicate the input prompt (decoder input ids) that are offset by 1. Now, we only override the -1
+ tokens in our prediction.
+ """
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input_ids.shape
+
+ max_length = max_length if max_length is not None else self.generation_config.max_length
+ input_ids_shifted = (
+ torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1
+ )
+
+ channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks
+ # we only apply the mask if we have a large enough seq len - otherwise we return as is
+ if max_length < 2 * channel_codebooks - 1:
+ return input_ids.reshape(bsz * num_codebooks, -1), input_ids_shifted.reshape(bsz * num_codebooks, -1)
+
+ # fill the shifted ids with the prompt entries, offset by the codebook idx
+ for codebook in range(channel_codebooks):
+ if self.config.audio_channels == 1:
+ # mono channel - loop over the codebooks one-by-one
+ input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook]
+ else:
+ # left/right channels are interleaved in the generated codebooks, so handle one then the other
+ input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook]
+ input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1]
+
+ # construct a pattern mask that indicates the positions of padding tokens for each codebook
+ # first fill the upper triangular part (the EOS padding)
+ delay_pattern = torch.triu(
+ torch.ones((channel_codebooks, max_length), dtype=torch.bool), diagonal=max_length - channel_codebooks + 1
+ )
+ # then fill the lower triangular part (the BOS padding)
+ delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.bool))
+
+ if self.config.audio_channels == 2:
+ # for left/right channel we need to duplicate every row of the pattern mask in an interleaved fashion
+ delay_pattern = delay_pattern.repeat_interleave(2, dim=0)
+
+ mask = ~delay_pattern.to(input_ids.device)
+ input_ids = mask * input_ids_shifted + ~mask * pad_token_id
+
+ # find the first position to start generating - this is the first place we have the -1 token
+ # and will always be in the first codebook (since it has no codebook offset)
+ first_codebook_ids = input_ids[:, 0, :]
+ start_ids = (first_codebook_ids == -1).nonzero()[:, 1]
+ if len(start_ids) > 0:
+ first_start_id = min(start_ids)
+ else:
+ # we have no tokens that need to be filled - return entire matrix of input ids
+ first_start_id = seq_len
+
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ pattern_mask = input_ids.reshape(bsz * num_codebooks, -1)
+ input_ids = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1)
+ return input_ids, pattern_mask
+
+ @staticmethod
+ def apply_delay_pattern_mask(input_ids, decoder_pad_token_mask):
+ """Apply a delay pattern mask to the decoder input ids, only preserving predictions where
+ the mask is set to -1, and otherwise setting to the value detailed in the mask."""
+ seq_len = input_ids.shape[-1]
+ decoder_pad_token_mask = decoder_pad_token_mask[..., :seq_len]
+ input_ids = torch.where(decoder_pad_token_mask == -1, input_ids, decoder_pad_token_mask)
+ return input_ids
+
+ @torch.no_grad()
+ # Ignore copy
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ synced_gpus: Optional[bool] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ if generation_config is None:
+ generation_config = self.generation_config
+
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+ generation_config.validate()
+ self._validate_model_kwargs(model_kwargs.copy())
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs`
+ input_ids, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = input_ids.shape[0] // self.num_codebooks
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=input_ids.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ input_ids, generation_config, model_kwargs
+ )
+
+ # 5. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=input_ids,
+ input_ids_length=input_ids_length,
+ )
+
+ # 6. Prepare `input_ids` which will be used for auto-regressive generation
+ # Build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to Musicgen)
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # stash the delay mask so that we don't have to recompute it in each forward pass
+ model_kwargs["delay_pattern_mask"] = delay_pattern_mask
+
+ # 7. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 8. prepare batched CFG externally (to enable coexistance with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 9. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=input_ids,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ **model_kwargs,
+ )
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=synced_gpus,
+ streamer=streamer,
+ **model_kwargs,
+ )
+
+ else:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.apply_delay_pattern_mask(output_ids, model_kwargs["delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.num_codebooks, -1
+ )
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_ids
+ return outputs
+ else:
+ return output_ids
+
+
+@add_start_docstrings(
+ "The composite Musicgen Melody model with a text and audio conditional models, a MusicgenMelody decoder and an audio encoder, "
+ "for music generation tasks with one or both of text and audio prompts.",
+ MUSICGEN_MELODY_START_DOCSTRING,
+ """
+ text_encoder (`Optional[PreTrainedModel]`, *optional*): Text encoder.
+ audio_encoder (`Optional[PreTrainedModel]`, *optional*): Audio code decoder.
+ decoder (`Optional[MusicgenMelodyForCausalLM]`, *optional*): MusicGen Melody decoder used to generate audio codes.
+ """,
+)
+class MusicgenMelodyForConditionalGeneration(PreTrainedModel, GenerationMixin):
+ config_class = MusicgenMelodyConfig
+ main_input_name = "input_ids"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ def __init__(
+ self,
+ config: MusicgenMelodyConfig = None,
+ text_encoder: Optional[PreTrainedModel] = None,
+ audio_encoder: Optional[PreTrainedModel] = None,
+ decoder: Optional[MusicgenMelodyForCausalLM] = None,
+ ):
+ if config is None and None in (text_encoder, audio_encoder, decoder):
+ raise ValueError(
+ "Either a configuration has to be provided, or all three of text encoder, audio encoder and Musicgen Melody decoder."
+ )
+ if config is None:
+ config = MusicgenMelodyConfig.from_sub_models_config(
+ text_encoder.config, audio_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}")
+
+ # initialize with config
+ super().__init__(config)
+
+ if text_encoder is None:
+ text_encoder = AutoModelForTextEncoding.from_config(config.text_encoder)
+
+ if audio_encoder is None:
+ audio_encoder = AutoModel.from_config(config.audio_encoder)
+
+ if decoder is None:
+ decoder = MusicgenMelodyForCausalLM._from_config(config.decoder)
+
+ self.text_encoder = text_encoder
+ self.audio_encoder = audio_encoder
+ self.decoder = 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.config.text_encoder._attn_implementation = self.text_encoder.config._attn_implementation
+ self.config.audio_encoder._attn_implementation = self.audio_encoder.config._attn_implementation
+ self.config.decoder._attn_implementation = self.decoder.config._attn_implementation
+ self.text_encoder.config = self.config.text_encoder
+ self.audio_encoder.config = self.config.audio_encoder
+ self.decoder.config = self.config.decoder
+
+ # text encoder outputs might need to be projected to different dimension for decoder
+ if self.text_encoder.config.hidden_size != self.decoder.config.hidden_size:
+ self.enc_to_dec_proj = nn.Linear(self.text_encoder.config.hidden_size, self.decoder.config.hidden_size)
+
+ # audio encoder outputs after chroma extraction might need to be projected to different dimension for decoder
+ if self.config.num_chroma != self.decoder.config.hidden_size:
+ self.audio_enc_to_dec_proj = nn.Linear(self.config.num_chroma, self.decoder.config.hidden_size)
+
+ if self.text_encoder.get_output_embeddings() is not None:
+ raise ValueError(
+ f"The encoder {self.text_encoder} should not have a LM Head. Please use a model without and LM Head"
+ )
+
+ # Initialize projection layers weights and tie text encoder and decoder weights if set accordingly
+ self.post_init()
+
+ def _init_weights(self, module):
+ # MusicgenMelodyForConditionalGeneration is made of PreTrainedModels that have already been initialized
+ # Projection layers still need to be initialized.
+ std = self.decoder.config.initializer_factor
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+
+ def tie_weights(self):
+ # tie text encoder & decoder if needed
+ if self.config.tie_encoder_decoder:
+ # tie text encoder and decoder base model
+ decoder_base_model_prefix = self.decoder.base_model_prefix
+ tied_weights = self._tie_encoder_decoder_weights(
+ self.text_encoder,
+ self.decoder._modules[decoder_base_model_prefix],
+ self.decoder.base_model_prefix,
+ "text_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_text_encoder(self):
+ return self.text_encoder
+
+ def get_encoder(self):
+ # get the text encoder to compute the conditionning hidden-states for generation
+ return self.get_text_encoder()
+
+ def get_decoder(self):
+ return self.decoder
+
+ def get_input_embeddings(self):
+ return self.text_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
+ # Copied from transformers.models.musicgen.modeling_musicgen.MusicgenForConditionalGeneration.from_sub_models_pretrained with Musicgen->MusicgenMelody, musicgen-small->musicgen-melody
+ def from_sub_models_pretrained(
+ cls,
+ text_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ audio_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ decoder_pretrained_model_name_or_path: Optional[str] = None,
+ *model_args,
+ **kwargs,
+ ) -> PreTrainedModel:
+ r"""
+ Instantiate a text encoder, an audio encoder, and a MusicGen decoder from one, two or three 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:
+ text_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text 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/`.
+
+ audio_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the audio 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/`.
+
+ 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/`.
+
+ 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 text encoder configuration, use the prefix *text_encoder_* for each configuration
+ parameter.
+ - To update the audio encoder configuration, use the prefix *audio_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 MusicgenMelodyForConditionalGeneration
+
+ >>> # initialize a musicgen model from a t5 text encoder, encodec audio encoder, and musicgen decoder
+ >>> model = MusicgenMelodyForConditionalGeneration.from_sub_models_pretrained(
+ ... text_encoder_pretrained_model_name_or_path="google-t5/t5-base",
+ ... audio_encoder_pretrained_model_name_or_path="facebook/encodec_24khz",
+ ... decoder_pretrained_model_name_or_path="facebook/musicgen-melody",
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./musicgen-ft")
+ >>> # load fine-tuned model
+ >>> model = MusicgenMelodyForConditionalGeneration.from_pretrained("./musicgen-ft")
+ ```"""
+
+ kwargs_text_encoder = {
+ argument[len("text_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("text_encoder_")
+ }
+
+ kwargs_audio_encoder = {
+ argument[len("audio_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("audio_encoder_")
+ }
+
+ kwargs_decoder = {
+ argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+ }
+
+ # remove text encoder, audio encoder and decoder kwargs from kwargs
+ for key in kwargs_text_encoder.keys():
+ del kwargs["text_encoder_" + key]
+ for key in kwargs_audio_encoder.keys():
+ del kwargs["audio_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.
+ text_encoder = kwargs_text_encoder.pop("model", None)
+ if text_encoder is None:
+ if text_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `text_encoder_model` is not defined as an argument, a `text_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_text_encoder:
+ encoder_config, kwargs_text_encoder = AutoConfig.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, **kwargs_text_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {text_encoder_pretrained_model_name_or_path} as a text_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_text_encoder["config"] = encoder_config
+
+ text_encoder = AutoModel.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, *model_args, **kwargs_text_encoder
+ )
+
+ audio_encoder = kwargs_audio_encoder.pop("model", None)
+ if audio_encoder is None:
+ if audio_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `audio_encoder_model` is not defined as an argument, an `audio_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_audio_encoder:
+ encoder_config, kwargs_audio_encoder = AutoConfig.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, **kwargs_audio_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {audio_encoder_pretrained_model_name_or_path} as an audio_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_audio_encoder["config"] = encoder_config
+
+ audio_encoder = AutoModel.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, *model_args, **kwargs_audio_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 isinstance(decoder_config, MusicgenMelodyConfig):
+ decoder_config = decoder_config.decoder
+
+ 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_sub_models_pretrained(...)` are set to `True` or do not pass a "
+ "`decoder_config` to `.from_sub_models_pretrained(...)`"
+ )
+
+ decoder = MusicgenMelodyForCausalLM.from_pretrained(
+ decoder_pretrained_model_name_or_path, **kwargs_decoder
+ )
+
+ # instantiate config with corresponding kwargs
+ config = MusicgenMelodyConfig.from_sub_models_config(
+ text_encoder.config, audio_encoder.config, decoder.config, **kwargs
+ )
+ return cls(text_encoder=text_encoder, audio_encoder=audio_encoder, decoder=decoder, config=config)
+
+ @add_start_docstrings_to_model_forward(MUSICGEN_MELODY_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=MusicgenMelodyOutputWithPast, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ input_features: Optional[torch.FloatTensor] = None,
+ decoder_input_ids: Optional[torch.LongTensor] = None,
+ decoder_attention_mask: Optional[torch.BoolTensor] = None,
+ past_key_values: Tuple[Tuple[torch.FloatTensor]] = None,
+ encoder_hidden_states: Optional[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, MusicgenMelodyOutputWithPast]:
+ r"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> from transformers import AutoProcessor, MusicgenMelodyForConditionalGeneration
+ >>> import torch
+
+ >>> processor = AutoProcessor.from_pretrained("facebook/musicgen-melody")
+ >>> model = MusicgenMelodyForConditionalGeneration.from_pretrained("facebook/musicgen-melody")
+
+ >>> inputs = processor(
+ ... text=["80s pop track with bassy drums and synth", "90s rock song with loud guitars and heavy drums"],
+ ... padding=True,
+ ... return_tensors="pt",
+ ... )
+
+ >>> pad_token_id = model.generation_config.pad_token_id
+ >>> decoder_input_ids = (
+ ... torch.ones((inputs.input_ids.shape[0] * model.decoder.num_codebooks, 1), dtype=torch.long)
+ ... * pad_token_id
+ ... )
+
+ >>> logits = model(**inputs, decoder_input_ids=decoder_input_ids).logits
+ >>> logits.shape # (bsz * num_codebooks, encoder_len + tgt_len, vocab_size)
+ torch.Size([8, 249, 2048])
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ kwargs_text_encoder = {
+ argument[len("text_encoder_")]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("text_encoder_")
+ }
+
+ kwargs_decoder = {
+ argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+ }
+
+ if encoder_hidden_states is None:
+ if inputs_embeds is not None or input_ids is not None:
+ encoder_outputs = self.text_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_text_encoder,
+ )
+
+ encoder_hidden_states = encoder_outputs[0]
+
+ # optionally project encoder_hidden_states
+ if self.text_encoder.config.hidden_size != self.decoder.config.hidden_size:
+ encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+
+ if attention_mask is not None and encoder_hidden_states is not None:
+ encoder_hidden_states = encoder_hidden_states * attention_mask[..., None]
+
+ # set a default audio conditional hidden states if text is not None
+ if encoder_hidden_states is not None and input_features is None:
+ input_features = torch.zeros(
+ (encoder_hidden_states.shape[0], 1, self.config.num_chroma),
+ device=self.device,
+ dtype=self.dtype,
+ )
+ input_features[:, :, 0] = 1
+
+ if input_features is not None:
+ audio_hidden_states = input_features
+
+ # optionally project audio_hidden_states ->
+ # (batch_size, seq_len, num_chroma) -> (batch_size, seq_len, hidden_size)
+ if self.config.num_chroma != self.decoder.config.hidden_size:
+ audio_hidden_states = self.audio_enc_to_dec_proj(audio_hidden_states)
+
+ # pad or truncate to config.chroma_length
+ if audio_hidden_states.shape[1] < self.config.chroma_length:
+ n_repeat = int(math.ceil(self.config.chroma_length / audio_hidden_states.shape[1]))
+ audio_hidden_states = audio_hidden_states.repeat(1, n_repeat, 1)
+ else:
+ logger.warning(
+ f"The conditional audio signal is of length {audio_hidden_states.shape[1]}, which exceeds"
+ f"the maximum chroma duration of {self.config.chroma_length}."
+ f"The audio will be truncated to {self.config.chroma_length} frames."
+ )
+ audio_hidden_states = audio_hidden_states[:, : self.config.chroma_length]
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states = torch.cat([audio_hidden_states, encoder_hidden_states], dim=1)
+ else:
+ encoder_hidden_states = audio_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.decoder.pad_token_id, self.config.decoder.bos_token_id
+ )
+
+ # Decode
+ decoder_outputs = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ 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,
+ labels=labels,
+ **kwargs_decoder,
+ )
+
+ if not return_dict:
+ return decoder_outputs + (encoder_hidden_states,)
+
+ return MusicgenMelodyOutputWithPast(
+ loss=decoder_outputs.loss,
+ logits=decoder_outputs.logits,
+ past_key_values=decoder_outputs.past_key_values,
+ hidden_states=decoder_outputs.hidden_states,
+ attentions=decoder_outputs.attentions,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ decoder_input_ids,
+ encoder_hidden_states=None,
+ past_key_values=None,
+ attention_mask=None,
+ decoder_attention_mask=None,
+ decoder_head_mask=None,
+ use_cache=None,
+ decoder_delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if decoder_delay_pattern_mask is None:
+ decoder_input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ decoder_input_ids,
+ self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ decoder_input_ids = self.decoder.apply_delay_pattern_mask(decoder_input_ids, decoder_delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ decoder_input_ids = decoder_input_ids.repeat((2, 1))
+ if decoder_attention_mask is not None:
+ decoder_attention_mask = decoder_attention_mask.repeat((2, 1))
+
+ 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 decoder_input_ids.shape[1] > past_length:
+ remove_prefix_length = past_length
+ else:
+ # Default to old behavior: keep only final ID
+ remove_prefix_length = decoder_input_ids.shape[1] - 1
+
+ decoder_input_ids = decoder_input_ids[:, remove_prefix_length:]
+
+ # we only want to use conditional signal in the 1st generation step but keeping the attention mask
+ encoder_hidden_states = None
+ # we also have to update the attention mask
+
+ return {
+ "input_ids": None, # encoder_hidden_states is defined. input_ids not needed
+ "encoder_hidden_states": encoder_hidden_states,
+ "past_key_values": past_key_values,
+ "decoder_input_ids": decoder_input_ids,
+ "attention_mask": attention_mask,
+ "decoder_attention_mask": decoder_attention_mask,
+ "decoder_head_mask": decoder_head_mask,
+ "use_cache": use_cache,
+ }
+
+ # Copied from transformers.models.musicgen.modeling_musicgen.MusicgenForConditionalGeneration._prepare_decoder_input_ids_for_generation
+ def _prepare_decoder_input_ids_for_generation(
+ self,
+ batch_size: int,
+ model_input_name: str,
+ model_kwargs: Dict[str, torch.Tensor],
+ decoder_start_token_id: Optional[int] = None,
+ bos_token_id: Optional[int] = None,
+ device: torch.device = None,
+ ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]:
+ """Prepares `decoder_input_ids` for generation with encoder-decoder models"""
+
+ # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming,
+ # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input.
+ if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
+ decoder_input_ids = model_kwargs.pop("decoder_input_ids")
+ elif "input_ids" in model_kwargs and model_input_name != "input_ids":
+ decoder_input_ids = model_kwargs.pop("input_ids")
+ else:
+ decoder_input_ids = None
+
+ # 2. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that.
+ decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)
+ if device is None:
+ device = self.device
+ decoder_input_ids_start = (
+ torch.ones((batch_size * self.decoder.num_codebooks, 1), dtype=torch.long, device=device)
+ * decoder_start_token_id
+ )
+
+ # no user input -> use decoder_start_token_id as decoder_input_ids
+ if decoder_input_ids is None:
+ decoder_input_ids = decoder_input_ids_start
+
+ # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust
+ # decoder_attention_mask if provided)
+ elif (decoder_input_ids[..., 0] != decoder_start_token_id).all().item():
+ decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
+ if "decoder_attention_mask" in model_kwargs:
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+ decoder_attention_mask = torch.cat(
+ (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),
+ dim=-1,
+ )
+ model_kwargs["decoder_attention_mask"] = decoder_attention_mask
+
+ return decoder_input_ids, model_kwargs
+
+ def _prepare_encoder_hidden_states_kwargs_for_generation(
+ self,
+ inputs_tensor: torch.Tensor,
+ model_kwargs,
+ model_input_name: Optional[str],
+ generation_config: GenerationConfig,
+ ) -> Dict[str, Any]:
+ encoder_hidden_states = None
+ # attention mask is consumed once to produce text conditional hidden states through the text encoder
+ encoder_attention_mask = model_kwargs.pop("attention_mask")
+ guidance_scale = generation_config.guidance_scale
+
+ # 1. condition on text
+ if inputs_tensor is not None:
+ encoder = self.get_text_encoder()
+ # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+ # as the inputs.
+ if hasattr(encoder, "_hf_hook"):
+ encoder._hf_hook.io_same_device = True
+
+ # Prepare args and kwargs from model kwargs.
+ irrelevant_prefix = ["decoder_", "use_cache"]
+ encoder_kwargs = {
+ argument: value
+ for argument, value in model_kwargs.items()
+ if not any(argument.startswith(p) for p in irrelevant_prefix)
+ }
+ encoder_signature = set(inspect.signature(encoder.forward).parameters)
+ encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+ if not encoder_accepts_wildcard:
+ encoder_kwargs = {
+ argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+ }
+ encoder_kwargs["output_attentions"] = generation_config.output_attentions
+ encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+
+ # make sure that encoder returns `ModelOutput`
+ model_input_name = model_input_name if model_input_name is not None else self.text_encoder.main_input_name
+ encoder_kwargs["return_dict"] = True
+ encoder_kwargs[model_input_name] = inputs_tensor
+ if encoder_attention_mask is not None:
+ encoder_kwargs["attention_mask"] = encoder_attention_mask
+ encoder_hidden_states = encoder(**encoder_kwargs).last_hidden_state
+
+ # optionally project encoder_hidden_states
+ if self.text_encoder.config.hidden_size != self.decoder.config.hidden_size:
+ encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+
+ # for classifier free guidance we need to add a 'null' input to our encoder hidden states
+ if guidance_scale is not None and guidance_scale > 1:
+ encoder_hidden_states = torch.concatenate(
+ [encoder_hidden_states, torch.zeros_like(encoder_hidden_states)], dim=0
+ )
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = torch.concatenate(
+ [encoder_attention_mask, torch.zeros_like(encoder_attention_mask)], dim=0
+ )
+ if encoder_attention_mask is not None:
+ encoder_hidden_states = encoder_hidden_states * encoder_attention_mask[..., None]
+
+ # 2. condition on audio
+ audio_hidden_states = model_kwargs.get("input_features", None)
+
+ if inputs_tensor is not None:
+ if audio_hidden_states is not None:
+ null_audio_hidden_states = torch.zeros_like(audio_hidden_states)
+ else:
+ null_audio_hidden_states = torch.zeros(
+ (inputs_tensor.shape[0], 1, self.config.num_chroma), device=self.device, dtype=self.dtype
+ )
+ null_audio_hidden_states[:, :, 0] = 1
+
+ if audio_hidden_states is None:
+ audio_hidden_states = null_audio_hidden_states
+
+ if audio_hidden_states is not None:
+ # for classifier free guidance we need to add a 'null' input to our audio hidden states
+ if guidance_scale is not None and guidance_scale > 1:
+ audio_hidden_states = torch.concatenate([audio_hidden_states, null_audio_hidden_states], dim=0)
+
+ # optionally project audio_hidden_states ->
+ # (batch_size, seq_len, num_chroma) -> (batch_size, seq_len, hidden_size)
+ if self.config.num_chroma != self.decoder.config.hidden_size:
+ audio_hidden_states = self.audio_enc_to_dec_proj(audio_hidden_states)
+
+ # pad or truncate to config.chroma_length
+ if audio_hidden_states.shape[1] < self.config.chroma_length:
+ n_repeat = int(math.ceil(self.config.chroma_length / audio_hidden_states.shape[1]))
+ audio_hidden_states = audio_hidden_states.repeat(1, n_repeat, 1)
+ audio_hidden_states = audio_hidden_states[:, : self.config.chroma_length]
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states = torch.cat([audio_hidden_states, encoder_hidden_states], dim=1)
+ else:
+ encoder_hidden_states = audio_hidden_states
+
+ model_kwargs["encoder_hidden_states"] = encoder_hidden_states
+
+ return model_kwargs
+
+ def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+ return shift_tokens_right(labels, self.config.decoder.pad_token_id, self.config.decoder.bos_token_id)
+
+ 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 _maybe_initialize_input_ids_for_generation(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ bos_token_id: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+ ) -> torch.LongTensor:
+ """Initializes input ids for generation, if necessary."""
+ if inputs is not None:
+ return inputs
+
+ if bos_token_id is None:
+ raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")
+
+ # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with
+ # soft-prompting or in multimodal implementations built on top of decoder-only language models.
+ batch_size = 1
+ for value in model_kwargs.values():
+ if isinstance(value, torch.Tensor):
+ batch_size = value.shape[0]
+ break
+ return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id
+
+ def freeze_audio_encoder(self):
+ """
+ Freeze the audio encoder weights.
+ """
+ for param in self.audio_encoder.parameters():
+ param.requires_grad = False
+ self.audio_encoder._requires_grad = False
+
+ def freeze_text_encoder(self):
+ """
+ Freeze the text encoder weights.
+ """
+ for param in self.text_encoder.parameters():
+ param.requires_grad = False
+ self.text_encoder._requires_grad = False
+
+ # Copied from transformers.models.musicgen.modeling_musicgen.MusicgenForConditionalGeneration._get_decoder_start_token_id
+ def _get_decoder_start_token_id(
+ self, decoder_start_token_id: Union[int, List[int]] = None, bos_token_id: Optional[int] = None
+ ) -> int:
+ decoder_start_token_id = (
+ decoder_start_token_id
+ if decoder_start_token_id is not None
+ else self.generation_config.decoder_start_token_id
+ )
+ bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
+
+ if decoder_start_token_id is not None:
+ return decoder_start_token_id
+ elif bos_token_id is not None:
+ return bos_token_id
+ raise ValueError(
+ "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ synced_gpus: Optional[bool] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ if generation_config is None:
+ generation_config = self.generation_config
+
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+ generation_config.validate()
+ self._validate_model_kwargs(model_kwargs.copy())
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs
+ inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = inputs_tensor.shape[0]
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=inputs_tensor.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ inputs_tensor, generation_config, model_kwargs
+ )
+
+ if "encoder_hidden_states" not in model_kwargs:
+ # encoder_hidden_states are created and added to `model_kwargs`
+ model_kwargs = self._prepare_encoder_hidden_states_kwargs_for_generation(
+ inputs_tensor, model_kwargs, model_input_name, generation_config
+ )
+
+ # 5. Prepare `input_ids` which will be used for auto-regressive generation
+ input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
+ batch_size=batch_size,
+ model_input_name=model_input_name,
+ model_kwargs=model_kwargs,
+ decoder_start_token_id=generation_config._decoder_start_token_tensor,
+ bos_token_id=generation_config._bos_token_tensor,
+ device=inputs_tensor.device,
+ )
+
+ # 6. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=inputs_tensor,
+ input_ids_length=input_ids_length,
+ )
+
+ # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+ input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+ # stash the delay mask so that we don't have to recompute in each forward pass
+ model_kwargs["decoder_delay_pattern_mask"] = decoder_delay_pattern_mask
+
+ # input_ids are ready to be placed on the streamer (if used)
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # 7. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 8. prepare batched CFG externally (to enable coexistance with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 9. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=inputs_tensor,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ is_encoder_decoder=self.config.is_encoder_decoder,
+ **model_kwargs,
+ )
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=synced_gpus,
+ streamer=streamer,
+ **model_kwargs,
+ )
+
+ else:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`."
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.decoder.apply_delay_pattern_mask(output_ids, model_kwargs["decoder_delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.decoder.num_codebooks, -1
+ )
+
+ # append the frame dimension back to the audio codes
+ output_ids = output_ids[None, ...]
+
+ audio_scales = model_kwargs.get("audio_scales")
+ if audio_scales is None:
+ audio_scales = [None] * batch_size
+
+ if self.decoder.config.audio_channels == 1:
+ output_values = self.audio_encoder.decode(
+ output_ids,
+ audio_scales=audio_scales,
+ ).audio_values
+ else:
+ codec_outputs_left = self.audio_encoder.decode(output_ids[:, :, ::2, :], audio_scales=audio_scales)
+ output_values_left = codec_outputs_left.audio_values
+
+ codec_outputs_right = self.audio_encoder.decode(output_ids[:, :, 1::2, :], audio_scales=audio_scales)
+ output_values_right = codec_outputs_right.audio_values
+
+ output_values = torch.cat([output_values_left, output_values_right], dim=1)
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_values
+ return outputs
+ else:
+ return output_values
diff --git a/lib/python3.12/site-packages/transformers/models/musicgen_melody/processing_musicgen_melody.py b/lib/python3.12/site-packages/transformers/models/musicgen_melody/processing_musicgen_melody.py
new file mode 100644
index 0000000000000000000000000000000000000000..8cf11e67d46d4d111cdf7a02ba045126060fc5f5
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/musicgen_melody/processing_musicgen_melody.py
@@ -0,0 +1,175 @@
+# coding=utf-8
+# Copyright 2024 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.
+"""
+Text/audio processor class for MusicGen Melody
+"""
+
+from typing import List, Optional
+
+import numpy as np
+
+from ...processing_utils import ProcessorMixin
+from ...utils import to_numpy
+
+
+class MusicgenMelodyProcessor(ProcessorMixin):
+ r"""
+ Constructs a MusicGen Melody processor which wraps a Wav2Vec2 feature extractor - for raw audio waveform processing - and a T5 tokenizer into a single processor
+ class.
+
+ [`MusicgenProcessor`] offers all the functionalities of [`MusicgenMelodyFeatureExtractor`] and [`T5Tokenizer`]. See
+ [`~MusicgenProcessor.__call__`] and [`~MusicgenProcessor.decode`] for more information.
+
+ Args:
+ feature_extractor (`MusicgenMelodyFeatureExtractor`):
+ An instance of [`MusicgenMelodyFeatureExtractor`]. The feature extractor is a required input.
+ tokenizer (`T5Tokenizer`):
+ An instance of [`T5Tokenizer`]. The tokenizer is a required input.
+ """
+
+ feature_extractor_class = "MusicgenMelodyFeatureExtractor"
+ tokenizer_class = ("T5Tokenizer", "T5TokenizerFast")
+
+ def __init__(self, feature_extractor, tokenizer):
+ super().__init__(feature_extractor, tokenizer)
+
+ # Copied from transformers.models.musicgen.processing_musicgen.MusicgenProcessor.get_decoder_prompt_ids
+ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
+ return self.tokenizer.get_decoder_prompt_ids(task=task, language=language, no_timestamps=no_timestamps)
+
+ def __call__(self, audio=None, text=None, **kwargs):
+ """
+ Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `audio`
+ and `kwargs` arguments to MusicgenMelodyFeatureExtractor's [`~MusicgenMelodyFeatureExtractor.__call__`] if `audio` is not
+ `None` to pre-process the audio. It also forwards the `text` and `kwargs` arguments to
+ PreTrainedTokenizer's [`~PreTrainedTokenizer.__call__`] if `text` is not `None`. Please refer to the docstring of the above two methods for more information.
+
+ Args:
+ audio (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The audio or batch of audios to be prepared. Each audio can be NumPy array or PyTorch tensor. In case
+ of a NumPy array/PyTorch tensor, each audio should be a mono-stereo signal of shape (T), where T is the sample length of the audio.
+ 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).
+ kwargs (*optional*):
+ Remaining dictionary of keyword arguments that will be passed to the feature extractor and/or the
+ tokenizer.
+ 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`.
+ - **input_features** -- Audio input features to be fed to a model. Returned when `audio` is not `None`.
+ - **attention_mask** -- List of token indices specifying which tokens should be attended to by the model when `text` is not `None`.
+ When only `audio` is specified, returns the timestamps attention mask.
+ """
+
+ sampling_rate = kwargs.pop("sampling_rate", None)
+
+ if audio is None and text is None:
+ raise ValueError("You need to specify either an `audio` or `text` input to process.")
+
+ if text is not None:
+ inputs = self.tokenizer(text, **kwargs)
+ if audio is not None:
+ audio_inputs = self.feature_extractor(audio, sampling_rate=sampling_rate, **kwargs)
+
+ if text is None:
+ return audio_inputs
+ elif audio is None:
+ return inputs
+ else:
+ inputs["input_features"] = audio_inputs["input_features"]
+ return inputs
+
+ # Copied from transformers.models.musicgen.processing_musicgen.MusicgenProcessor.batch_decode with padding_mask->attention_mask
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method is used to decode either batches of audio outputs from the MusicGen model, or batches of token ids
+ from the tokenizer. In the case of decoding token ids, this method forwards all its arguments to T5Tokenizer's
+ [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information.
+ """
+ audio_values = kwargs.pop("audio", None)
+ attention_mask = kwargs.pop("attention_mask", None)
+
+ if len(args) > 0:
+ audio_values = args[0]
+ args = args[1:]
+
+ if audio_values is not None:
+ return self._decode_audio(audio_values, attention_mask=attention_mask)
+ else:
+ return self.tokenizer.batch_decode(*args, **kwargs)
+
+ # Copied from transformers.models.musicgen.processing_musicgen.MusicgenProcessor.decode
+ def decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to T5Tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to the
+ docstring of this method for more information.
+ """
+ return self.tokenizer.decode(*args, **kwargs)
+
+ # Copied from transformers.models.musicgen.processing_musicgen.MusicgenProcessor._decode_audio with padding_mask->attention_mask
+ def _decode_audio(self, audio_values, attention_mask: Optional = None) -> List[np.ndarray]:
+ """
+ This method strips any padding from the audio values to return a list of numpy audio arrays.
+ """
+ audio_values = to_numpy(audio_values)
+ bsz, channels, seq_len = audio_values.shape
+
+ if attention_mask is None:
+ return list(audio_values)
+
+ attention_mask = to_numpy(attention_mask)
+
+ # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
+ # token (so that the generated audio values are **not** treated as padded tokens)
+ difference = seq_len - attention_mask.shape[-1]
+ padding_value = 1 - self.feature_extractor.padding_value
+ attention_mask = np.pad(attention_mask, ((0, 0), (0, difference)), "constant", constant_values=padding_value)
+
+ audio_values = audio_values.tolist()
+ for i in range(bsz):
+ sliced_audio = np.asarray(audio_values[i])[
+ attention_mask[i][None, :] != self.feature_extractor.padding_value
+ ]
+ audio_values[i] = sliced_audio.reshape(channels, -1)
+
+ return audio_values
+
+ def get_unconditional_inputs(self, num_samples=1, return_tensors="pt"):
+ """
+ Helper function to get null inputs for unconditional generation, enabling the model to be used without the
+ feature extractor or tokenizer.
+
+ Args:
+ num_samples (int, *optional*):
+ Number of audio samples to unconditionally generate.
+
+ Example:
+ ```python
+ >>> from transformers import MusicgenMelodyForConditionalGeneration, MusicgenMelodyProcessor
+
+ >>> model = MusicgenMelodyForConditionalGeneration.from_pretrained("facebook/musicgen-melody")
+
+ >>> # get the unconditional (or 'null') inputs for the model
+ >>> processor = MusicgenMelodyProcessor.from_pretrained("facebook/musicgen-melody")
+ >>> unconditional_inputs = processor.get_unconditional_inputs(num_samples=1)
+
+ >>> audio_samples = model.generate(**unconditional_inputs, max_new_tokens=256)
+ ```"""
+ inputs = self.tokenizer([""] * num_samples, return_tensors=return_tensors, return_attention_mask=True)
+ inputs["attention_mask"][:] = 0
+
+ return inputs
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/__init__.py b/lib/python3.12/site-packages/transformers/models/nllb/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5cdb326098a30020f0d9d2cc4c083090bce70d5a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/nllb/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .tokenization_nllb import *
+ from .tokenization_nllb_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bac91cb29475429c8276f60da6b927cf78ac3eac
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3b0c0834e39408eee9584b5fde5473d2374440b2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5551f61a5d81f39587dfdac8f252a5ad677c318c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/nllb/__pycache__/tokenization_nllb_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb.py b/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb.py
new file mode 100644
index 0000000000000000000000000000000000000000..02e07f0ecadd3f8df8ed5a4c6be6d055dd0a7f39
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb.py
@@ -0,0 +1,392 @@
+# coding=utf-8
+# Copyright 2022 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.
+
+import os
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+
+from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+SPIECE_UNDERLINE = "▁"
+
+VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.bpe.model"}
+
+
+FAIRSEQ_LANGUAGE_CODES = ['ace_Arab', 'ace_Latn', 'acm_Arab', 'acq_Arab', 'aeb_Arab', 'afr_Latn', 'ajp_Arab', 'aka_Latn', 'amh_Ethi', 'apc_Arab', 'arb_Arab', 'ars_Arab', 'ary_Arab', 'arz_Arab', 'asm_Beng', 'ast_Latn', 'awa_Deva', 'ayr_Latn', 'azb_Arab', 'azj_Latn', 'bak_Cyrl', 'bam_Latn', 'ban_Latn', 'bel_Cyrl', 'bem_Latn', 'ben_Beng', 'bho_Deva', 'bjn_Arab', 'bjn_Latn', 'bod_Tibt', 'bos_Latn', 'bug_Latn', 'bul_Cyrl', 'cat_Latn', 'ceb_Latn', 'ces_Latn', 'cjk_Latn', 'ckb_Arab', 'crh_Latn', 'cym_Latn', 'dan_Latn', 'deu_Latn', 'dik_Latn', 'dyu_Latn', 'dzo_Tibt', 'ell_Grek', 'eng_Latn', 'epo_Latn', 'est_Latn', 'eus_Latn', 'ewe_Latn', 'fao_Latn', 'pes_Arab', 'fij_Latn', 'fin_Latn', 'fon_Latn', 'fra_Latn', 'fur_Latn', 'fuv_Latn', 'gla_Latn', 'gle_Latn', 'glg_Latn', 'grn_Latn', 'guj_Gujr', 'hat_Latn', 'hau_Latn', 'heb_Hebr', 'hin_Deva', 'hne_Deva', 'hrv_Latn', 'hun_Latn', 'hye_Armn', 'ibo_Latn', 'ilo_Latn', 'ind_Latn', 'isl_Latn', 'ita_Latn', 'jav_Latn', 'jpn_Jpan', 'kab_Latn', 'kac_Latn', 'kam_Latn', 'kan_Knda', 'kas_Arab', 'kas_Deva', 'kat_Geor', 'knc_Arab', 'knc_Latn', 'kaz_Cyrl', 'kbp_Latn', 'kea_Latn', 'khm_Khmr', 'kik_Latn', 'kin_Latn', 'kir_Cyrl', 'kmb_Latn', 'kon_Latn', 'kor_Hang', 'kmr_Latn', 'lao_Laoo', 'lvs_Latn', 'lij_Latn', 'lim_Latn', 'lin_Latn', 'lit_Latn', 'lmo_Latn', 'ltg_Latn', 'ltz_Latn', 'lua_Latn', 'lug_Latn', 'luo_Latn', 'lus_Latn', 'mag_Deva', 'mai_Deva', 'mal_Mlym', 'mar_Deva', 'min_Latn', 'mkd_Cyrl', 'plt_Latn', 'mlt_Latn', 'mni_Beng', 'khk_Cyrl', 'mos_Latn', 'mri_Latn', 'zsm_Latn', 'mya_Mymr', 'nld_Latn', 'nno_Latn', 'nob_Latn', 'npi_Deva', 'nso_Latn', 'nus_Latn', 'nya_Latn', 'oci_Latn', 'gaz_Latn', 'ory_Orya', 'pag_Latn', 'pan_Guru', 'pap_Latn', 'pol_Latn', 'por_Latn', 'prs_Arab', 'pbt_Arab', 'quy_Latn', 'ron_Latn', 'run_Latn', 'rus_Cyrl', 'sag_Latn', 'san_Deva', 'sat_Beng', 'scn_Latn', 'shn_Mymr', 'sin_Sinh', 'slk_Latn', 'slv_Latn', 'smo_Latn', 'sna_Latn', 'snd_Arab', 'som_Latn', 'sot_Latn', 'spa_Latn', 'als_Latn', 'srd_Latn', 'srp_Cyrl', 'ssw_Latn', 'sun_Latn', 'swe_Latn', 'swh_Latn', 'szl_Latn', 'tam_Taml', 'tat_Cyrl', 'tel_Telu', 'tgk_Cyrl', 'tgl_Latn', 'tha_Thai', 'tir_Ethi', 'taq_Latn', 'taq_Tfng', 'tpi_Latn', 'tsn_Latn', 'tso_Latn', 'tuk_Latn', 'tum_Latn', 'tur_Latn', 'twi_Latn', 'tzm_Tfng', 'uig_Arab', 'ukr_Cyrl', 'umb_Latn', 'urd_Arab', 'uzn_Latn', 'vec_Latn', 'vie_Latn', 'war_Latn', 'wol_Latn', 'xho_Latn', 'ydd_Hebr', 'yor_Latn', 'yue_Hant', 'zho_Hans', 'zho_Hant', 'zul_Latn'] # fmt: skip
+
+
+class NllbTokenizer(PreTrainedTokenizer):
+ """
+ Construct an NLLB tokenizer.
+
+ Adapted from [`RobertaTokenizer`] and [`XLNetTokenizer`]. Based on
+ [SentencePiece](https://github.com/google/sentencepiece).
+
+ The tokenization method is ` ` for source language documents, and `
+ ` for target language documents.
+
+ Examples:
+
+ ```python
+ >>> from transformers import NllbTokenizer
+
+ >>> tokenizer = NllbTokenizer.from_pretrained(
+ ... "facebook/nllb-200-distilled-600M", src_lang="eng_Latn", tgt_lang="fra_Latn"
+ ... )
+ >>> example_english_phrase = " UN Chief Says There Is No Military Solution in Syria"
+ >>> expected_translation_french = "Le chef de l'ONU affirme qu'il n'y a pas de solution militaire en Syrie."
+ >>> inputs = tokenizer(example_english_phrase, text_target=expected_translation_french, return_tensors="pt")
+ ```
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ 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.
+ tokenizer_file (`str`, *optional*):
+ The path to a tokenizer file to use instead of the vocab file.
+ src_lang (`str`, *optional*):
+ The language to use as source language for translation.
+ tgt_lang (`str`, *optional*):
+ The language to use as target language for translation.
+ sp_model_kwargs (`Dict[str, str]`):
+ Additional keyword arguments to pass to the model initialization.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+
+ prefix_tokens: List[int] = []
+ suffix_tokens: List[int] = []
+
+ def __init__(
+ self,
+ vocab_file,
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ tokenizer_file=None,
+ src_lang=None,
+ tgt_lang=None,
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
+ additional_special_tokens=None,
+ legacy_behaviour=False,
+ **kwargs,
+ ):
+ if additional_special_tokens is None:
+ additional_special_tokens = FAIRSEQ_LANGUAGE_CODES
+ bos_token = AddedToken(bos_token, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
+ pad_token = AddedToken(pad_token, normalized=False, special=True) if isinstance(pad_token, str) else pad_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
+ # Mask token behave like a normal word, i.e. include the space before it
+ mask_token = (
+ AddedToken(mask_token, normalized=True, lstrip=True, 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.legacy_behaviour = legacy_behaviour
+
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+ self.sp_model.Load(str(vocab_file))
+ self.vocab_file = vocab_file
+ # Original fairseq vocab and spm vocab must be "aligned":
+ # Vocab | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9
+ # -------- | ------- | ------- | ------ | ------- | ---- | ---- | ---- | ---- | ---- | ----
+ # fairseq | '' | '' | '' | '' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a'
+ # spm | '' | '' | '' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a' | '▁s'
+
+ # unk token needs to be in the vocab with correct index
+ self._added_tokens_decoder = {0: bos_token, 1: pad_token, 2: eos_token, 3: unk_token}
+ # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
+ self.fairseq_offset = 1
+ self.sp_model_size = len(self.sp_model)
+
+ super().__init__(
+ 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,
+ tokenizer_file=tokenizer_file,
+ src_lang=src_lang,
+ tgt_lang=tgt_lang,
+ additional_special_tokens=additional_special_tokens,
+ sp_model_kwargs=self.sp_model_kwargs,
+ legacy_behaviour=legacy_behaviour,
+ **kwargs,
+ )
+
+ self._src_lang = src_lang if src_lang is not None else "eng_Latn"
+ self.cur_lang_code_id = self.convert_tokens_to_ids(self._src_lang)
+ self.tgt_lang = tgt_lang
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ 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
+
+ # 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.LoadFromSerializedProto(self.sp_model_proto)
+
+ @property
+ def vocab_size(self):
+ return len(self.sp_model) + self.fairseq_offset
+
+ @property
+ def src_lang(self) -> str:
+ return self._src_lang
+
+ @src_lang.setter
+ def src_lang(self, new_src_lang: str) -> None:
+ self._src_lang = new_src_lang
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ 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
+ )
+
+ prefix_ones = [1] * len(self.prefix_tokens)
+ suffix_ones = [1] * len(self.suffix_tokens)
+ if token_ids_1 is None:
+ return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
+ return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
+
+ 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 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.prefix_tokens + token_ids_0 + self.suffix_tokens
+ # We don't expect to process pairs, but leave the pair logic for API consistency
+ return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
+
+ 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. nllb does not
+ make use of token type ids, therefore a list of zeros is returned.
+
+ 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 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 _build_translation_inputs(
+ self, raw_inputs, return_tensors: str, src_lang: Optional[str], tgt_lang: Optional[str], **extra_kwargs
+ ):
+ """Used by translation pipeline, to prepare inputs for the generate function"""
+ if src_lang is None or tgt_lang is None:
+ raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model")
+ self.src_lang = src_lang
+ inputs = self(raw_inputs, add_special_tokens=True, return_tensors=return_tensors, **extra_kwargs)
+ tgt_lang_id = self.convert_tokens_to_ids(tgt_lang)
+ inputs["forced_bos_token_id"] = tgt_lang_id
+ return inputs
+
+ 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 _tokenize(self, text: str) -> List[str]:
+ return self.sp_model.encode(text, out_type=str)
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ spm_id = self.sp_model.PieceToId(token)
+ # Need to return unknown token if the SP model returned 0
+ return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.sp_model.IdToPiece(index - self.fairseq_offset)
+
+ 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 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 prepare_seq2seq_batch(
+ self,
+ src_texts: List[str],
+ src_lang: str = "eng_Latn",
+ tgt_texts: Optional[List[str]] = None,
+ tgt_lang: str = "fra_Latn",
+ **kwargs,
+ ) -> BatchEncoding:
+ self.src_lang = src_lang
+ self.tgt_lang = tgt_lang
+ return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
+
+ def _switch_to_input_mode(self):
+ return self.set_src_lang_special_tokens(self.src_lang)
+
+ def _switch_to_target_mode(self):
+ return self.set_tgt_lang_special_tokens(self.tgt_lang)
+
+ def set_src_lang_special_tokens(self, src_lang) -> None:
+ """Reset the special tokens to the source lang setting.
+ - In legacy mode: No prefix and suffix=[eos, src_lang_code].
+ - In default mode: Prefix=[src_lang_code], suffix = [eos]
+ """
+ self.cur_lang_code = self.convert_tokens_to_ids(src_lang)
+ if self.legacy_behaviour:
+ self.prefix_tokens = []
+ self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+ else:
+ self.prefix_tokens = [self.cur_lang_code]
+ self.suffix_tokens = [self.eos_token_id]
+
+ def set_tgt_lang_special_tokens(self, lang: str) -> None:
+ """Reset the special tokens to the target lang setting.
+ - In legacy mode: No prefix and suffix=[eos, tgt_lang_code].
+ - In default mode: Prefix=[tgt_lang_code], suffix = [eos]
+ """
+ self.cur_lang_code = self.convert_tokens_to_ids(lang)
+ if self.legacy_behaviour:
+ self.prefix_tokens = []
+ self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+ else:
+ self.prefix_tokens = [self.cur_lang_code]
+ self.suffix_tokens = [self.eos_token_id]
+
+
+__all__ = ["NllbTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb_fast.py b/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..80b00e342462d1b15e7d8e75964800a8e9365468
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/nllb/tokenization_nllb_fast.py
@@ -0,0 +1,331 @@
+# coding=utf-8
+# Copyright 2022 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.
+
+import os
+from shutil import copyfile
+from typing import List, Optional, Tuple
+
+from tokenizers import processors
+
+from ...tokenization_utils import AddedToken, BatchEncoding
+from ...tokenization_utils_fast import PreTrainedTokenizerFast
+from ...utils import is_sentencepiece_available, logging
+
+
+if is_sentencepiece_available():
+ from .tokenization_nllb import NllbTokenizer
+else:
+ NllbTokenizer = None
+
+
+logger = logging.get_logger(__name__)
+
+
+VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"}
+
+
+FAIRSEQ_LANGUAGE_CODES = ['ace_Arab', 'ace_Latn', 'acm_Arab', 'acq_Arab', 'aeb_Arab', 'afr_Latn', 'ajp_Arab', 'aka_Latn', 'amh_Ethi', 'apc_Arab', 'arb_Arab', 'ars_Arab', 'ary_Arab', 'arz_Arab', 'asm_Beng', 'ast_Latn', 'awa_Deva', 'ayr_Latn', 'azb_Arab', 'azj_Latn', 'bak_Cyrl', 'bam_Latn', 'ban_Latn', 'bel_Cyrl', 'bem_Latn', 'ben_Beng', 'bho_Deva', 'bjn_Arab', 'bjn_Latn', 'bod_Tibt', 'bos_Latn', 'bug_Latn', 'bul_Cyrl', 'cat_Latn', 'ceb_Latn', 'ces_Latn', 'cjk_Latn', 'ckb_Arab', 'crh_Latn', 'cym_Latn', 'dan_Latn', 'deu_Latn', 'dik_Latn', 'dyu_Latn', 'dzo_Tibt', 'ell_Grek', 'eng_Latn', 'epo_Latn', 'est_Latn', 'eus_Latn', 'ewe_Latn', 'fao_Latn', 'pes_Arab', 'fij_Latn', 'fin_Latn', 'fon_Latn', 'fra_Latn', 'fur_Latn', 'fuv_Latn', 'gla_Latn', 'gle_Latn', 'glg_Latn', 'grn_Latn', 'guj_Gujr', 'hat_Latn', 'hau_Latn', 'heb_Hebr', 'hin_Deva', 'hne_Deva', 'hrv_Latn', 'hun_Latn', 'hye_Armn', 'ibo_Latn', 'ilo_Latn', 'ind_Latn', 'isl_Latn', 'ita_Latn', 'jav_Latn', 'jpn_Jpan', 'kab_Latn', 'kac_Latn', 'kam_Latn', 'kan_Knda', 'kas_Arab', 'kas_Deva', 'kat_Geor', 'knc_Arab', 'knc_Latn', 'kaz_Cyrl', 'kbp_Latn', 'kea_Latn', 'khm_Khmr', 'kik_Latn', 'kin_Latn', 'kir_Cyrl', 'kmb_Latn', 'kon_Latn', 'kor_Hang', 'kmr_Latn', 'lao_Laoo', 'lvs_Latn', 'lij_Latn', 'lim_Latn', 'lin_Latn', 'lit_Latn', 'lmo_Latn', 'ltg_Latn', 'ltz_Latn', 'lua_Latn', 'lug_Latn', 'luo_Latn', 'lus_Latn', 'mag_Deva', 'mai_Deva', 'mal_Mlym', 'mar_Deva', 'min_Latn', 'mkd_Cyrl', 'plt_Latn', 'mlt_Latn', 'mni_Beng', 'khk_Cyrl', 'mos_Latn', 'mri_Latn', 'zsm_Latn', 'mya_Mymr', 'nld_Latn', 'nno_Latn', 'nob_Latn', 'npi_Deva', 'nso_Latn', 'nus_Latn', 'nya_Latn', 'oci_Latn', 'gaz_Latn', 'ory_Orya', 'pag_Latn', 'pan_Guru', 'pap_Latn', 'pol_Latn', 'por_Latn', 'prs_Arab', 'pbt_Arab', 'quy_Latn', 'ron_Latn', 'run_Latn', 'rus_Cyrl', 'sag_Latn', 'san_Deva', 'sat_Beng', 'scn_Latn', 'shn_Mymr', 'sin_Sinh', 'slk_Latn', 'slv_Latn', 'smo_Latn', 'sna_Latn', 'snd_Arab', 'som_Latn', 'sot_Latn', 'spa_Latn', 'als_Latn', 'srd_Latn', 'srp_Cyrl', 'ssw_Latn', 'sun_Latn', 'swe_Latn', 'swh_Latn', 'szl_Latn', 'tam_Taml', 'tat_Cyrl', 'tel_Telu', 'tgk_Cyrl', 'tgl_Latn', 'tha_Thai', 'tir_Ethi', 'taq_Latn', 'taq_Tfng', 'tpi_Latn', 'tsn_Latn', 'tso_Latn', 'tuk_Latn', 'tum_Latn', 'tur_Latn', 'twi_Latn', 'tzm_Tfng', 'uig_Arab', 'ukr_Cyrl', 'umb_Latn', 'urd_Arab', 'uzn_Latn', 'vec_Latn', 'vie_Latn', 'war_Latn', 'wol_Latn', 'xho_Latn', 'ydd_Hebr', 'yor_Latn', 'yue_Hant', 'zho_Hans', 'zho_Hant', 'zul_Latn'] # fmt: skip
+
+
+class NllbTokenizerFast(PreTrainedTokenizerFast):
+ """
+ Construct a "fast" NLLB tokenizer (backed by HuggingFace's *tokenizers* library). Based on
+ [BPE](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=BPE#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.
+
+ The tokenization method is ` ` for source language documents, and `
+ ` for target language documents.
+
+ Examples:
+
+ ```python
+ >>> from transformers import NllbTokenizerFast
+
+ >>> tokenizer = NllbTokenizerFast.from_pretrained(
+ ... "facebook/nllb-200-distilled-600M", src_lang="eng_Latn", tgt_lang="fra_Latn"
+ ... )
+ >>> example_english_phrase = " UN Chief Says There Is No Military Solution in Syria"
+ >>> expected_translation_french = "Le chef de l'ONU affirme qu'il n'y a pas de solution militaire en Syrie."
+ >>> inputs = tokenizer(example_english_phrase, text_target=expected_translation_french, return_tensors="pt")
+ ```
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ 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.
+ tokenizer_file (`str`, *optional*):
+ The path to a tokenizer file to use instead of the vocab file.
+ src_lang (`str`, *optional*):
+ The language to use as source language for translation.
+ tgt_lang (`str`, *optional*):
+ The language to use as target language for translation.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ slow_tokenizer_class = NllbTokenizer
+
+ prefix_tokens: List[int] = []
+ suffix_tokens: List[int] = []
+
+ def __init__(
+ self,
+ vocab_file=None,
+ tokenizer_file=None,
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ src_lang=None,
+ tgt_lang=None,
+ additional_special_tokens=None,
+ legacy_behaviour=False,
+ **kwargs,
+ ):
+ if additional_special_tokens is None:
+ additional_special_tokens = FAIRSEQ_LANGUAGE_CODES
+
+ self.vocab_file = vocab_file
+ # Mask token behave like a normal word, i.e. include the space before it
+ mask_token = (
+ AddedToken(mask_token, normalized=True, lstrip=True, special=True)
+ if isinstance(mask_token, str)
+ else mask_token
+ )
+ self.legacy_behaviour = legacy_behaviour
+ super().__init__(
+ vocab_file=vocab_file,
+ tokenizer_file=tokenizer_file,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ cls_token=cls_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ src_lang=src_lang,
+ tgt_lang=tgt_lang,
+ mask_token=mask_token,
+ additional_special_tokens=additional_special_tokens,
+ legacy_behaviour=legacy_behaviour,
+ **kwargs,
+ )
+
+ self._src_lang = src_lang if src_lang is not None else "eng_Latn"
+ self.cur_lang_code = self.convert_tokens_to_ids(self._src_lang)
+ self.tgt_lang = tgt_lang
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ @property
+ def can_save_slow_tokenizer(self) -> bool:
+ return os.path.isfile(self.vocab_file) if self.vocab_file else False
+
+ @property
+ def src_lang(self) -> str:
+ return self._src_lang
+
+ @src_lang.setter
+ def src_lang(self, new_src_lang: str) -> None:
+ self._src_lang = new_src_lang
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ 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.prefix_tokens + token_ids_0 + self.suffix_tokens
+ # We don't expect to process pairs, but leave the pair logic for API consistency
+ return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
+
+ 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. nllb does not
+ make use of token type ids, therefore a list of zeros is returned.
+
+ 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 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 _build_translation_inputs(
+ self, raw_inputs, return_tensors: str, src_lang: Optional[str], tgt_lang: Optional[str], **extra_kwargs
+ ):
+ """Used by translation pipeline, to prepare inputs for the generate function"""
+ if src_lang is None or tgt_lang is None:
+ raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model")
+ self.src_lang = src_lang
+ inputs = self(raw_inputs, add_special_tokens=True, return_tensors=return_tensors, **extra_kwargs)
+ tgt_lang_id = self.convert_tokens_to_ids(tgt_lang)
+ inputs["forced_bos_token_id"] = tgt_lang_id
+ return inputs
+
+ def prepare_seq2seq_batch(
+ self,
+ src_texts: List[str],
+ src_lang: str = "eng_Latn",
+ tgt_texts: Optional[List[str]] = None,
+ tgt_lang: str = "fra_Latn",
+ **kwargs,
+ ) -> BatchEncoding:
+ self.src_lang = src_lang
+ self.tgt_lang = tgt_lang
+ return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
+
+ def _switch_to_input_mode(self):
+ return self.set_src_lang_special_tokens(self.src_lang)
+
+ def _switch_to_target_mode(self):
+ return self.set_tgt_lang_special_tokens(self.tgt_lang)
+
+ def set_src_lang_special_tokens(self, src_lang) -> None:
+ """Reset the special tokens to the source lang setting.
+ - In legacy mode: No prefix and suffix=[eos, src_lang_code].
+ - In default mode: Prefix=[src_lang_code], suffix = [eos]
+ """
+ self.cur_lang_code = self.convert_tokens_to_ids(src_lang)
+
+ if self.legacy_behaviour:
+ self.prefix_tokens = []
+ self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+ else:
+ self.prefix_tokens = [self.cur_lang_code]
+ self.suffix_tokens = [self.eos_token_id]
+
+ prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens)
+ suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens)
+
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single=prefix_tokens_str + ["$A"] + suffix_tokens_str,
+ pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str,
+ special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str, self.prefix_tokens + self.suffix_tokens)),
+ )
+
+ def set_tgt_lang_special_tokens(self, lang: str) -> None:
+ """Reset the special tokens to the target lang setting.
+ - In legacy mode: No prefix and suffix=[eos, tgt_lang_code].
+ - In default mode: Prefix=[tgt_lang_code], suffix = [eos]
+ """
+ self.cur_lang_code = self.convert_tokens_to_ids(lang)
+ if self.legacy_behaviour:
+ self.prefix_tokens = []
+ self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+ else:
+ self.prefix_tokens = [self.cur_lang_code]
+ self.suffix_tokens = [self.eos_token_id]
+
+ prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens)
+ suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens)
+
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single=prefix_tokens_str + ["$A"] + suffix_tokens_str,
+ pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str,
+ special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str, self.prefix_tokens + self.suffix_tokens)),
+ )
+
+ 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,)
+
+
+__all__ = ["NllbTokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__init__.py b/lib/python3.12/site-packages/transformers/models/owlv2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ce329420882f7b02546a85addfca96dc19d35aa
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/owlv2/__init__.py
@@ -0,0 +1,29 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_owlv2 import *
+ from .image_processing_owlv2 import *
+ from .modeling_owlv2 import *
+ from .processing_owlv2 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b27e5f429c61459858483cefd9c1bbc9df94ad90
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/configuration_owlv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/configuration_owlv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..031329417380800af426a96f17e06ee7a0b4ecd2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/configuration_owlv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/image_processing_owlv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/image_processing_owlv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e5842d9a18303e5a2eb2c8d43cbba0804fa62e7c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/image_processing_owlv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/modeling_owlv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/modeling_owlv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..158fe43f1ed01f3b22020ae785ef8dbe01109473
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/modeling_owlv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/processing_owlv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/processing_owlv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..02d7357b849b6236a9ecb143dd88db4543c5de67
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/owlv2/__pycache__/processing_owlv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/configuration_owlv2.py b/lib/python3.12/site-packages/transformers/models/owlv2/configuration_owlv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf70afc7ccf527cc87bf8aa1f0b0016799c8bf9b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/owlv2/configuration_owlv2.py
@@ -0,0 +1,289 @@
+# 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.
+"""OWLv2 model configuration"""
+
+from typing import TYPE_CHECKING, Dict
+
+
+if TYPE_CHECKING:
+ pass
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.owlvit.configuration_owlvit.OwlViTTextConfig with OwlViT->Owlv2, owlvit-base-patch32->owlv2-base-patch16, owlvit->owlv2, OWL-ViT->OWLv2
+class Owlv2TextConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of an [`Owlv2TextModel`]. It is used to instantiate an
+ Owlv2 text 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 Owlv2
+ [google/owlv2-base-patch16](https://huggingface.co/google/owlv2-base-patch16) 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 49408):
+ Vocabulary size of the OWLv2 text model. Defines the number of different tokens that can be represented
+ by the `inputs_ids` passed when calling [`Owlv2TextModel`].
+ hidden_size (`int`, *optional*, defaults to 512):
+ Dimensionality of the encoder layers and the pooler layer.
+ intermediate_size (`int`, *optional*, defaults to 2048):
+ 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 8):
+ Number of attention heads for each attention layer in the Transformer encoder.
+ max_position_embeddings (`int`, *optional*, defaults to 16):
+ 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_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-05):
+ 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.
+ initializer_factor (`float`, *optional*, defaults to 1.0):
+ A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+ testing).
+ pad_token_id (`int`, *optional*, defaults to 0):
+ The id of the padding token in the input sequences.
+ bos_token_id (`int`, *optional*, defaults to 49406):
+ The id of the beginning-of-sequence token in the input sequences.
+ eos_token_id (`int`, *optional*, defaults to 49407):
+ The id of the end-of-sequence token in the input sequences.
+
+ Example:
+
+ ```python
+ >>> from transformers import Owlv2TextConfig, Owlv2TextModel
+
+ >>> # Initializing a Owlv2TextModel with google/owlv2-base-patch16 style configuration
+ >>> configuration = Owlv2TextConfig()
+
+ >>> # Initializing a Owlv2TextConfig from the google/owlv2-base-patch16 style configuration
+ >>> model = Owlv2TextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "owlv2_text_model"
+ base_config_key = "text_config"
+
+ def __init__(
+ self,
+ vocab_size=49408,
+ hidden_size=512,
+ intermediate_size=2048,
+ num_hidden_layers=12,
+ num_attention_heads=8,
+ max_position_embeddings=16,
+ hidden_act="quick_gelu",
+ layer_norm_eps=1e-5,
+ attention_dropout=0.0,
+ initializer_range=0.02,
+ initializer_factor=1.0,
+ pad_token_id=0,
+ bos_token_id=49406,
+ eos_token_id=49407,
+ **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.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.attention_dropout = attention_dropout
+ self.initializer_range = initializer_range
+ self.initializer_factor = initializer_factor
+
+
+# Copied from transformers.models.owlvit.configuration_owlvit.OwlViTVisionConfig with OwlViT->Owlv2, owlvit-base-patch32->owlv2-base-patch16, owlvit->owlv2, OWL-ViT->OWLv2, 32->16
+class Owlv2VisionConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of an [`Owlv2VisionModel`]. It is used to instantiate
+ an OWLv2 image 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 OWLv2
+ [google/owlv2-base-patch16](https://huggingface.co/google/owlv2-base-patch16) 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.
+ num_channels (`int`, *optional*, defaults to 3):
+ Number of channels in the input images.
+ image_size (`int`, *optional*, defaults to 768):
+ 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-05):
+ 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.
+ initializer_factor (`float`, *optional*, defaults to 1.0):
+ A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+ testing).
+
+ Example:
+
+ ```python
+ >>> from transformers import Owlv2VisionConfig, Owlv2VisionModel
+
+ >>> # Initializing a Owlv2VisionModel with google/owlv2-base-patch16 style configuration
+ >>> configuration = Owlv2VisionConfig()
+
+ >>> # Initializing a Owlv2VisionModel model from the google/owlv2-base-patch16 style configuration
+ >>> model = Owlv2VisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "owlv2_vision_model"
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ hidden_size=768,
+ intermediate_size=3072,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ num_channels=3,
+ image_size=768,
+ patch_size=16,
+ hidden_act="quick_gelu",
+ layer_norm_eps=1e-5,
+ attention_dropout=0.0,
+ initializer_range=0.02,
+ initializer_factor=1.0,
+ **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.image_size = image_size
+ self.patch_size = patch_size
+ self.hidden_act = hidden_act
+ self.layer_norm_eps = layer_norm_eps
+ self.attention_dropout = attention_dropout
+ self.initializer_range = initializer_range
+ self.initializer_factor = initializer_factor
+
+
+# Copied from transformers.models.owlvit.configuration_owlvit.OwlViTConfig with OwlViT->Owlv2, owlvit-base-patch32->owlv2-base-patch16, owlvit->owlv2, OWL-ViT->OWLv2
+class Owlv2Config(PretrainedConfig):
+ r"""
+ [`Owlv2Config`] is the configuration class to store the configuration of an [`Owlv2Model`]. It is used to
+ instantiate an OWLv2 model according to the specified arguments, defining the text model and vision model
+ configs. Instantiating a configuration with the defaults will yield a similar configuration to that of the OWLv2
+ [google/owlv2-base-patch16](https://huggingface.co/google/owlv2-base-patch16) 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 [`Owlv2TextConfig`].
+ vision_config (`dict`, *optional*):
+ Dictionary of configuration options used to initialize [`Owlv2VisionConfig`].
+ projection_dim (`int`, *optional*, defaults to 512):
+ Dimensionality of text and vision projection layers.
+ logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
+ The initial value of the *logit_scale* parameter. Default is used as per the original OWLv2
+ implementation.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return a dictionary. If `False`, returns a tuple.
+ kwargs (*optional*):
+ Dictionary of keyword arguments.
+ """
+
+ model_type = "owlv2"
+ sub_configs = {"text_config": Owlv2TextConfig, "vision_config": Owlv2VisionConfig}
+
+ def __init__(
+ self,
+ text_config=None,
+ vision_config=None,
+ projection_dim=512,
+ logit_scale_init_value=2.6592,
+ return_dict=True,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ if text_config is None:
+ text_config = {}
+ logger.info("text_config is None. Initializing the Owlv2TextConfig with default values.")
+
+ if vision_config is None:
+ vision_config = {}
+ logger.info("vision_config is None. initializing the Owlv2VisionConfig with default values.")
+
+ self.text_config = Owlv2TextConfig(**text_config)
+ self.vision_config = Owlv2VisionConfig(**vision_config)
+
+ self.projection_dim = projection_dim
+ self.logit_scale_init_value = logit_scale_init_value
+ self.return_dict = return_dict
+ self.initializer_factor = 1.0
+
+ @classmethod
+ def from_text_vision_configs(cls, text_config: Dict, vision_config: Dict, **kwargs):
+ r"""
+ Instantiate a [`Owlv2Config`] (or a derived class) from owlv2 text model configuration and owlv2 vision
+ model configuration.
+
+ Returns:
+ [`Owlv2Config`]: An instance of a configuration object
+ """
+ config_dict = {}
+ config_dict["text_config"] = text_config
+ config_dict["vision_config"] = vision_config
+
+ return cls.from_dict(config_dict, **kwargs)
+
+
+__all__ = ["Owlv2Config", "Owlv2TextConfig", "Owlv2VisionConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/image_processing_owlv2.py b/lib/python3.12/site-packages/transformers/models/owlv2/image_processing_owlv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..56c8075ef0ee085dd81412284d9bb50d65f2b15b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/owlv2/image_processing_owlv2.py
@@ -0,0 +1,635 @@
+# 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 OWLv2."""
+
+import warnings
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from ...image_transforms import (
+ center_to_corners_format,
+ pad,
+ to_channel_dimension_format,
+)
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ get_image_size,
+ infer_channel_dimension_format,
+ is_scaled_image,
+ make_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import (
+ TensorType,
+ filter_out_non_signature_kwargs,
+ is_scipy_available,
+ is_torch_available,
+ is_vision_available,
+ logging,
+ requires_backends,
+)
+
+
+if is_torch_available():
+ import torch
+
+
+if is_vision_available():
+ import PIL
+
+if is_scipy_available():
+ from scipy import ndimage as ndi
+
+if TYPE_CHECKING:
+ from .modeling_owlv2 import Owlv2ObjectDetectionOutput
+
+logger = logging.get_logger(__name__)
+
+
+def _scale_boxes(boxes, target_sizes):
+ """
+ Scale batch of bounding boxes to the target sizes.
+
+ Args:
+ boxes (`torch.Tensor` of shape `(batch_size, num_boxes, 4)`):
+ Bounding boxes to scale. Each box is expected to be in (x1, y1, x2, y2) format.
+ target_sizes (`List[Tuple[int, int]]` or `torch.Tensor` of shape `(batch_size, 2)`):
+ Target sizes to scale the boxes to. Each target size is expected to be in (height, width) format.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size, num_boxes, 4)`: Scaled bounding boxes.
+ """
+
+ if isinstance(target_sizes, (list, tuple)):
+ image_height = torch.tensor([i[0] for i in target_sizes])
+ image_width = torch.tensor([i[1] for i in target_sizes])
+ elif isinstance(target_sizes, torch.Tensor):
+ image_height, image_width = target_sizes.unbind(1)
+ else:
+ raise ValueError("`target_sizes` must be a list, tuple or torch.Tensor")
+
+ # for owlv2 image is padded to max size unlike owlvit, thats why we have to scale boxes to max size
+ max_size = torch.max(image_height, image_width)
+
+ scale_factor = torch.stack([max_size, max_size, max_size, max_size], dim=1)
+ scale_factor = scale_factor.unsqueeze(1).to(boxes.device)
+ boxes = boxes * scale_factor
+ return boxes
+
+
+# Copied from transformers.models.owlvit.image_processing_owlvit._upcast
+def _upcast(t):
+ # 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.owlvit.image_processing_owlvit.box_area
+def box_area(boxes):
+ """
+ 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.owlvit.image_processing_owlvit.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
+
+
+def _preprocess_resize_output_shape(image, output_shape):
+ """Validate resize output shape according to input image.
+
+ Args:
+ image (`np.ndarray`):
+ Image to be resized.
+ output_shape (`iterable`):
+ Size of the generated output image `(rows, cols[, ...][, dim])`. If `dim` is not provided, the number of
+ channels is preserved.
+
+ Returns
+ image (`np.ndarray`):
+ The input image, but with additional singleton dimensions appended in the case where `len(output_shape) >
+ input.ndim`.
+ output_shape (`Tuple`):
+ The output shape converted to tuple.
+
+ Raises ------ ValueError:
+ If output_shape length is smaller than the image number of dimensions.
+
+ Notes ----- The input image is reshaped if its number of dimensions is not equal to output_shape_length.
+
+ """
+ output_shape = tuple(output_shape)
+ output_ndim = len(output_shape)
+ input_shape = image.shape
+ if output_ndim > image.ndim:
+ # append dimensions to input_shape
+ input_shape += (1,) * (output_ndim - image.ndim)
+ image = np.reshape(image, input_shape)
+ elif output_ndim == image.ndim - 1:
+ # multichannel case: append shape of last axis
+ output_shape = output_shape + (image.shape[-1],)
+ elif output_ndim < image.ndim:
+ raise ValueError("output_shape length cannot be smaller than the image number of dimensions")
+
+ return image, output_shape
+
+
+def _clip_warp_output(input_image, output_image):
+ """Clip output image to range of values of input image.
+
+ Note that this function modifies the values of *output_image* in-place.
+
+ Taken from:
+ https://github.com/scikit-image/scikit-image/blob/b4b521d6f0a105aabeaa31699949f78453ca3511/skimage/transform/_warps.py#L640.
+
+ Args:
+ input_image : ndarray
+ Input image.
+ output_image : ndarray
+ Output image, which is modified in-place.
+ """
+ min_val = np.min(input_image)
+ if np.isnan(min_val):
+ # NaNs detected, use NaN-safe min/max
+ min_func = np.nanmin
+ max_func = np.nanmax
+ min_val = min_func(input_image)
+ else:
+ min_func = np.min
+ max_func = np.max
+ max_val = max_func(input_image)
+
+ output_image = np.clip(output_image, min_val, max_val)
+
+ return output_image
+
+
+class Owlv2ImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs an OWLv2 image processor.
+
+ Args:
+ 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_pad (`bool`, *optional*, defaults to `True`):
+ Whether to pad the image to a square with gray pixels on the bottom and the right. Can be overriden by
+ `do_pad` in the `preprocess` method.
+ 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 `{"height": 960, "width": 960}`):
+ Size to resize the image to. Can be overriden by `size` in the `preprocess` method.
+ resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
+ Resampling method to use if resizing the image. Can be overriden by `resample` 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 `OPENAI_CLIP_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 `OPENAI_CLIP_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_rescale: bool = True,
+ rescale_factor: Union[int, float] = 1 / 255,
+ do_pad: bool = True,
+ do_resize: bool = True,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = PILImageResampling.BILINEAR,
+ 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)
+
+ self.do_rescale = do_rescale
+ self.rescale_factor = rescale_factor
+ self.do_pad = do_pad
+ self.do_resize = do_resize
+ self.size = size if size is not None else {"height": 960, "width": 960}
+ self.resample = resample
+ self.do_normalize = do_normalize
+ self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+ self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+
+ def pad(
+ self,
+ image: np.array,
+ data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ):
+ """
+ Pad an image to a square with gray pixels on the bottom and the right, as per the original OWLv2
+ implementation.
+
+ Args:
+ image (`np.ndarray`):
+ Image to pad.
+ 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.
+ """
+ height, width = get_image_size(image)
+ size = max(height, width)
+ image = pad(
+ image=image,
+ padding=((0, size - height), (0, size - width)),
+ constant_values=0.5,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+
+ return image
+
+ def resize(
+ self,
+ image: np.ndarray,
+ size: Dict[str, int],
+ anti_aliasing: bool = True,
+ anti_aliasing_sigma=None,
+ data_format: Optional[Union[str, ChannelDimension]] = None,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ **kwargs,
+ ) -> np.ndarray:
+ """
+ Resize an image as per the original implementation.
+
+ Args:
+ image (`np.ndarray`):
+ Image to resize.
+ size (`Dict[str, int]`):
+ Dictionary containing the height and width to resize the image to.
+ anti_aliasing (`bool`, *optional*, defaults to `True`):
+ Whether to apply anti-aliasing when downsampling the image.
+ anti_aliasing_sigma (`float`, *optional*, defaults to `None`):
+ Standard deviation for Gaussian kernel when downsampling the image. If `None`, it will be calculated
+ automatically.
+ 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.
+ """
+ requires_backends(self, "scipy")
+
+ output_shape = (size["height"], size["width"])
+ image = to_channel_dimension_format(image, ChannelDimension.LAST)
+ image, output_shape = _preprocess_resize_output_shape(image, output_shape)
+ input_shape = image.shape
+ factors = np.divide(input_shape, output_shape)
+
+ # Translate modes used by np.pad to those used by scipy.ndimage
+ ndi_mode = "mirror"
+ cval = 0
+ order = 1
+ if anti_aliasing:
+ if anti_aliasing_sigma is None:
+ anti_aliasing_sigma = np.maximum(0, (factors - 1) / 2)
+ else:
+ anti_aliasing_sigma = np.atleast_1d(anti_aliasing_sigma) * np.ones_like(factors)
+ if np.any(anti_aliasing_sigma < 0):
+ raise ValueError("Anti-aliasing standard deviation must be greater than or equal to zero")
+ elif np.any((anti_aliasing_sigma > 0) & (factors <= 1)):
+ warnings.warn(
+ "Anti-aliasing standard deviation greater than zero but not down-sampling along all axes"
+ )
+ filtered = ndi.gaussian_filter(image, anti_aliasing_sigma, cval=cval, mode=ndi_mode)
+ else:
+ filtered = image
+
+ zoom_factors = [1 / f for f in factors]
+ out = ndi.zoom(filtered, zoom_factors, order=order, mode=ndi_mode, cval=cval, grid_mode=True)
+
+ image = _clip_warp_output(image, out)
+
+ image = to_channel_dimension_format(image, input_data_format, ChannelDimension.LAST)
+ image = (
+ to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image
+ )
+ return image
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_pad: Optional[bool] = None,
+ do_resize: Optional[bool] = None,
+ size: 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,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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_pad (`bool`, *optional*, defaults to `self.do_pad`):
+ Whether to pad the image to a square with gray pixels on the bottom and the right.
+ do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+ Whether to resize the image.
+ size (`Dict[str, int]`, *optional*, defaults to `self.size`):
+ Size to resize the image to.
+ 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_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_pad = do_pad if do_pad is not None else self.do_pad
+ do_resize = do_resize if do_resize is not None else self.do_resize
+ 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) # for BC
+
+ 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."
+ )
+ # Here, pad and resize methods are different from the rest of image processors
+ # as they don't have any resampling in resize()
+ # or pad size in pad() (the maximum of (height, width) is taken instead).
+ # hence, these arguments don't need to be passed in validate_preprocess_arguments.
+ validate_preprocess_arguments(
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ size=size,
+ )
+
+ # All transformations expect numpy arrays.
+ images = [to_numpy_array(image) for image in images]
+
+ if do_rescale and is_scaled_image(images[0]):
+ 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_rescale:
+ images = [
+ self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
+ for image in images
+ ]
+
+ if do_pad:
+ images = [self.pad(image=image, input_data_format=input_data_format) for image in images]
+
+ if do_resize:
+ images = [
+ self.resize(
+ image=image,
+ size=size,
+ 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)
+
+ # Copied from transformers.models.owlvit.image_processing_owlvit.OwlViTImageProcessor.post_process_object_detection with OwlViT->Owlv2
+ def post_process_object_detection(
+ self,
+ outputs: "Owlv2ObjectDetectionOutput",
+ threshold: float = 0.1,
+ target_sizes: Optional[Union[TensorType, List[Tuple]]] = None,
+ ):
+ """
+ Converts the raw output of [`Owlv2ForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
+ bottom_right_x, bottom_right_y) format.
+
+ Args:
+ outputs ([`Owlv2ObjectDetectionOutput`]):
+ Raw outputs of the model.
+ threshold (`float`, *optional*, defaults to 0.1):
+ 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 following keys:
+ - "scores": The confidence scores for each predicted box on the image.
+ - "labels": Indexes of the classes predicted by the model on the image.
+ - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.
+ """
+ batch_logits, batch_boxes = outputs.logits, outputs.pred_boxes
+ batch_size = len(batch_logits)
+
+ if target_sizes is not None and len(target_sizes) != batch_size:
+ raise ValueError("Make sure that you pass in as many target sizes as images")
+
+ # batch_logits of shape (batch_size, num_queries, num_classes)
+ batch_class_logits = torch.max(batch_logits, dim=-1)
+ batch_scores = torch.sigmoid(batch_class_logits.values)
+ batch_labels = batch_class_logits.indices
+
+ # Convert to [x0, y0, x1, y1] format
+ batch_boxes = center_to_corners_format(batch_boxes)
+
+ # Convert from relative [0, 1] to absolute [0, height] coordinates
+ if target_sizes is not None:
+ batch_boxes = _scale_boxes(batch_boxes, target_sizes)
+
+ results = []
+ for scores, labels, boxes in zip(batch_scores, batch_labels, batch_boxes):
+ keep = scores > threshold
+ scores = scores[keep]
+ labels = labels[keep]
+ boxes = boxes[keep]
+ results.append({"scores": scores, "labels": labels, "boxes": boxes})
+
+ return results
+
+ # Copied from transformers.models.owlvit.image_processing_owlvit.OwlViTImageProcessor.post_process_image_guided_detection
+ def post_process_image_guided_detection(self, outputs, threshold=0.0, nms_threshold=0.3, target_sizes=None):
+ """
+ Converts the output of [`OwlViTForObjectDetection.image_guided_detection`] into the format expected by the COCO
+ api.
+
+ Args:
+ outputs ([`OwlViTImageGuidedObjectDetectionOutput`]):
+ Raw outputs of the model.
+ threshold (`float`, *optional*, defaults to 0.0):
+ Minimum confidence threshold to use to filter out predicted boxes.
+ nms_threshold (`float`, *optional*, defaults to 0.3):
+ IoU threshold for non-maximum suppression of overlapping boxes.
+ 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, predicted normalized bounding boxes are rescaled to the target sizes. If left to
+ None, predictions will not be unnormalized.
+
+ 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. All labels are set to None as
+ `OwlViTForObjectDetection.image_guided_detection` perform one-shot object detection.
+ """
+ logits, target_boxes = outputs.logits, outputs.target_pred_boxes
+
+ if target_sizes is not None and 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 target_sizes is not None and target_sizes.shape[1] != 2:
+ raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
+
+ probs = torch.max(logits, dim=-1)
+ scores = torch.sigmoid(probs.values)
+
+ # Convert to [x0, y0, x1, y1] format
+ target_boxes = center_to_corners_format(target_boxes)
+
+ # Apply non-maximum suppression (NMS)
+ if nms_threshold < 1.0:
+ for idx in range(target_boxes.shape[0]):
+ for i in torch.argsort(-scores[idx]):
+ if not scores[idx][i]:
+ continue
+
+ ious = box_iou(target_boxes[idx][i, :].unsqueeze(0), target_boxes[idx])[0][0]
+ ious[i] = -1.0 # Mask self-IoU.
+ scores[idx][ious > nms_threshold] = 0.0
+
+ # Convert from relative [0, 1] to absolute [0, height] coordinates
+ if target_sizes is not None:
+ target_boxes = _scale_boxes(target_boxes, target_sizes)
+
+ # Compute box display alphas based on prediction scores
+ results = []
+ alphas = torch.zeros_like(scores)
+
+ for idx in range(target_boxes.shape[0]):
+ # Select scores for boxes matching the current query:
+ query_scores = scores[idx]
+ if not query_scores.nonzero().numel():
+ continue
+
+ # Apply threshold on scores before scaling
+ query_scores[query_scores < threshold] = 0.0
+
+ # Scale box alpha such that the best box for each query has alpha 1.0 and the worst box has alpha 0.1.
+ # All other boxes will either belong to a different query, or will not be shown.
+ max_score = torch.max(query_scores) + 1e-6
+ query_alphas = (query_scores - (max_score * 0.1)) / (max_score * 0.9)
+ query_alphas = torch.clip(query_alphas, 0.0, 1.0)
+ alphas[idx] = query_alphas
+
+ mask = alphas[idx] > 0
+ box_scores = alphas[idx][mask]
+ boxes = target_boxes[idx][mask]
+ results.append({"scores": box_scores, "labels": None, "boxes": boxes})
+
+ return results
+
+
+__all__ = ["Owlv2ImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/modeling_owlv2.py b/lib/python3.12/site-packages/transformers/models/owlv2/modeling_owlv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d83ab0004d878e7488d636c0f0d55cb9ede6f79
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/owlv2/modeling_owlv2.py
@@ -0,0 +1,1845 @@
+# coding=utf-8
+# Copyright 2023 Google AI 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 OWLv2 model."""
+
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import Tensor, nn
+
+from ...activations import ACT2FN
+from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ ModelOutput,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_vision_available,
+ logging,
+ replace_return_docstrings,
+ torch_int,
+)
+from .configuration_owlv2 import Owlv2Config, Owlv2TextConfig, Owlv2VisionConfig
+
+
+if is_vision_available():
+ from transformers.image_transforms import center_to_corners_format
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "google/owlv2-base-patch16-ensemble"
+
+# See all Owlv2 models at https://huggingface.co/models?filter=owlv2
+
+
+# Copied from transformers.models.clip.modeling_clip.contrastive_loss with clip->owlv2
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+# Copied from transformers.models.clip.modeling_clip.clip_loss with clip->owlv2
+def owlv2_loss(similarity: torch.Tensor) -> torch.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(similarity.t())
+ return (caption_loss + image_loss) / 2.0
+
+
+@dataclass
+class Owlv2Output(ModelOutput):
+ """
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Contrastive loss for image-text similarity.
+ logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+ The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
+ similarity scores.
+ logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+ The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
+ similarity scores.
+ text_embeds (`torch.FloatTensor` of shape `(batch_size * num_max_text_queries, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of [`Owlv2TextModel`].
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The image embeddings obtained by applying the projection layer to the pooled output of
+ [`Owlv2VisionModel`].
+ text_model_output (Tuple[`BaseModelOutputWithPooling`]):
+ The output of the [`Owlv2TextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`Owlv2VisionModel`].
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits_per_image: Optional[torch.FloatTensor] = None
+ logits_per_text: Optional[torch.FloatTensor] = None
+ text_embeds: Optional[torch.FloatTensor] = None
+ image_embeds: Optional[torch.FloatTensor] = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> Tuple[Any]:
+ return tuple(
+ self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+# Copied from transformers.loss.loss_for_object_detection._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.loss.loss_for_object_detection.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.loss.loss_for_object_detection.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.loss.loss_for_object_detection.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
+
+
+@dataclass
+class Owlv2ObjectDetectionOutput(ModelOutput):
+ """
+ Output type of [`Owlv2ForObjectDetection`].
+
+ 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_patches, num_queries)`):
+ Classification logits (including no-object) for all queries.
+ objectness_logits (`torch.FloatTensor` of shape `(batch_size, num_patches, 1)`):
+ The objectness logits of all image patches. OWL-ViT represents images as a set of image patches where the
+ total number of patches is (image_size / patch_size)**2.
+ pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 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 [`~Owlv2ImageProcessor.post_process_object_detection`] to retrieve the
+ unnormalized bounding boxes.
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, num_max_text_queries, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of [`Owlv2TextModel`].
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):
+ Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes image
+ embeddings for each patch.
+ class_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`):
+ Class embeddings of all image patches. OWLv2 represents images as a set of image patches where the total
+ number of patches is (image_size / patch_size)**2.
+ text_model_output (Tuple[`BaseModelOutputWithPooling`]):
+ The output of the [`Owlv2TextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`Owlv2VisionModel`].
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ loss_dict: Optional[Dict] = None
+ logits: Optional[torch.FloatTensor] = None
+ objectness_logits: Optional[torch.FloatTensor] = None
+ pred_boxes: Optional[torch.FloatTensor] = None
+ text_embeds: Optional[torch.FloatTensor] = None
+ image_embeds: Optional[torch.FloatTensor] = None
+ class_embeds: Optional[torch.FloatTensor] = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> Tuple[Any]:
+ return tuple(
+ self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+@dataclass
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTImageGuidedObjectDetectionOutput with OwlViT->Owlv2,OWL-ViT->OWLv2
+class Owlv2ImageGuidedObjectDetectionOutput(ModelOutput):
+ """
+ Output type of [`Owlv2ForObjectDetection.image_guided_detection`].
+
+ Args:
+ logits (`torch.FloatTensor` of shape `(batch_size, num_patches, num_queries)`):
+ Classification logits (including no-object) for all queries.
+ target_pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 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 target image in the batch
+ (disregarding possible padding). You can use [`~Owlv2ImageProcessor.post_process_object_detection`] to
+ retrieve the unnormalized bounding boxes.
+ query_pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 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 query image in the batch
+ (disregarding possible padding). You can use [`~Owlv2ImageProcessor.post_process_object_detection`] to
+ retrieve the unnormalized bounding boxes.
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):
+ Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes
+ image embeddings for each patch.
+ query_image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):
+ Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes
+ image embeddings for each patch.
+ class_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`):
+ Class embeddings of all image patches. OWLv2 represents images as a set of image patches where the total
+ number of patches is (image_size / patch_size)**2.
+ text_model_output (Tuple[`BaseModelOutputWithPooling`]):
+ The output of the [`Owlv2TextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`Owlv2VisionModel`].
+ """
+
+ logits: Optional[torch.FloatTensor] = None
+ image_embeds: Optional[torch.FloatTensor] = None
+ query_image_embeds: Optional[torch.FloatTensor] = None
+ target_pred_boxes: Optional[torch.FloatTensor] = None
+ query_pred_boxes: Optional[torch.FloatTensor] = None
+ class_embeds: Optional[torch.FloatTensor] = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> Tuple[Any]:
+ return tuple(
+ self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionEmbeddings with OwlViT->Owlv2
+class Owlv2VisionEmbeddings(nn.Module):
+ def __init__(self, config: Owlv2VisionConfig):
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.class_embedding = nn.Parameter(torch.randn(config.hidden_size))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=config.patch_size,
+ stride=config.patch_size,
+ bias=False,
+ )
+
+ self.num_patches = (config.image_size // config.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)
+
+ # Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings.interpolate_pos_encoding
+ 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. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ position_embedding = self.position_embedding.weight.unsqueeze(0)
+ num_positions = position_embedding.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embedding(self.position_ids)
+
+ class_pos_embed = position_embedding[:, :1]
+ patch_pos_embed = position_embedding[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, _, height, width = pixel_values.shape
+ patch_embeds = self.patch_embedding(pixel_values) # shape = [batch_size, num_channels, height, width]
+ 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)
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+ return embeddings
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextEmbeddings with OwlViT->Owlv2
+class Owlv2TextEmbeddings(nn.Module):
+ def __init__(self, config: Owlv2TextConfig):
+ super().__init__()
+ self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.position_embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+ # 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: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ ) -> torch.Tensor:
+ seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+
+ if inputs_embeds is None:
+ inputs_embeds = self.token_embedding(input_ids)
+
+ position_embeddings = self.position_embedding(position_ids)
+ embeddings = inputs_embeds + position_embeddings
+
+ return embeddings
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTAttention with OwlViT->Owlv2
+class Owlv2Attention(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)
+
+ # For int8 compatibility, sometimes the `attn_probs` are in `fp32`
+ attn_probs = attn_probs.to(value_states.dtype)
+
+ 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.CLIPMLP with CLIP->Owlv2
+class Owlv2MLP(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.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->Owlv2
+class Owlv2EncoderLayer(nn.Module):
+ def __init__(self, config: Owlv2Config):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = Owlv2Attention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = Owlv2MLP(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.owlvit.modeling_owlvit.OwlViTPreTrainedModel with OwlViT->Owlv2,owlvit->owlv2
+class Owlv2PreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = Owlv2Config
+ base_model_prefix = "owlv2"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Owlv2EncoderLayer"]
+
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ factor = self.config.initializer_factor
+ if isinstance(module, Owlv2TextEmbeddings):
+ module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
+ module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
+ elif isinstance(module, Owlv2VisionEmbeddings):
+ factor = self.config.initializer_factor
+ nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
+ nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
+ nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
+ elif isinstance(module, Owlv2Attention):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ out_proj_std = (module.embed_dim**-0.5) * factor
+ nn.init.normal_(module.q_proj.weight, std=in_proj_std)
+ nn.init.normal_(module.k_proj.weight, std=in_proj_std)
+ nn.init.normal_(module.v_proj.weight, std=in_proj_std)
+ nn.init.normal_(module.out_proj.weight, std=out_proj_std)
+ elif isinstance(module, Owlv2MLP):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
+ nn.init.normal_(module.fc1.weight, std=fc_std)
+ nn.init.normal_(module.fc2.weight, std=in_proj_std)
+ elif isinstance(module, Owlv2Model):
+ nn.init.normal_(
+ module.text_projection.weight,
+ std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+ nn.init.normal_(
+ module.visual_projection.weight,
+ std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+ if isinstance(module, nn.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+ if isinstance(module, nn.Linear) and module.bias is not None:
+ module.bias.data.zero_()
+
+
+OWLV2_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 ([`Owvl2Config`]): 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.
+"""
+
+OWLV2_TEXT_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, 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.Tensor` of shape `(batch_size, num_max_text_queries, 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)
+ 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.
+"""
+
+OWLV2_VISION_INPUTS_DOCSTRING = r"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel 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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+OWLV2_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)
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values.
+ 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)
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+ 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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_base_image_embeds (`bool`, *optional*):
+ Whether or not to return the base image embeddings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+OWLV2_OBJECT_DETECTION_INPUTS_DOCSTRING = r"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values.
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`, *optional*):
+ 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.Tensor` of shape `(batch_size, num_max_text_queries, 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)
+ output_hidden_states (`bool`, *optional*):
+ Whether or not to return the last hidden state. See `text_model_last_hidden_state` and
+ `vision_model_last_hidden_state` under returned tensors for more detail.
+ interpolate_pos_encoding (`bool`, *optional*, defaults `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+OWLV2_IMAGE_GUIDED_OBJECT_DETECTION_INPUTS_DOCSTRING = r"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values.
+ query_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values of query image(s) to be detected. Pass in one query image per target image.
+ 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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTEncoder with OwlViT->Owlv2
+class Owlv2Encoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`Owlv2EncoderLayer`].
+
+ Args:
+ config: Owlv2Config
+ """
+
+ def __init__(self, config: Owlv2Config):
+ super().__init__()
+ self.layers = nn.ModuleList([Owlv2EncoderLayer(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)`).
+ 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 encoder_layer in 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
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextTransformer with OWLVIT->OWLV2,OwlViT->Owlv2
+class Owlv2TextTransformer(nn.Module):
+ def __init__(self, config: Owlv2TextConfig):
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+ self.embeddings = Owlv2TextEmbeddings(config)
+ self.encoder = Owlv2Encoder(config)
+ self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ @add_start_docstrings_to_model_forward(OWLV2_TEXT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Owlv2TextConfig)
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: 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,
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ 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
+
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+ hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+ # num_samples, seq_len = input_shape where num_samples = batch_size * num_max_text_queries
+ # OWLV2's text model uses causal mask, prepare it here.
+ # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+ causal_attention_mask = _create_4d_causal_attention_mask(
+ input_shape, hidden_states.dtype, device=hidden_states.device
+ )
+ # expand attention_mask
+ if attention_mask is not None:
+ # [num_samples, seq_len] -> [num_samples, 1, tgt_seq_len, src_seq_len]
+ attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+ encoder_outputs = self.encoder(
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ causal_attention_mask=causal_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+ last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+ # take features from the end of tokens embedding (end of token is the highest number in each sequence)
+ # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+ pooled_output = last_hidden_state[
+ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+ input_ids.to(torch.int).argmax(dim=-1).to(last_hidden_state.device),
+ ]
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextModel with google/owlvit-base-patch32->google/owlv2-base-patch16, OWLVIT->OWLV2,OwlViT->Owlv2
+class Owlv2TextModel(Owlv2PreTrainedModel):
+ config_class = Owlv2TextConfig
+
+ def __init__(self, config: Owlv2TextConfig):
+ super().__init__(config)
+ self.text_model = Owlv2TextTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.text_model.embeddings.token_embedding
+
+ def set_input_embeddings(self, value):
+ self.text_model.embeddings.token_embedding = value
+
+ @add_start_docstrings_to_model_forward(OWLV2_TEXT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Owlv2TextConfig)
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ attention_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"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> from transformers import AutoProcessor, Owlv2TextModel
+
+ >>> model = Owlv2TextModel.from_pretrained("google/owlv2-base-patch16")
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16")
+ >>> inputs = processor(
+ ... text=[["a photo of a cat", "a photo of a dog"], ["photo of a astranaut"]], return_tensors="pt"
+ ... )
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled (EOS token) states
+ ```"""
+
+ # Get embeddings for all text queries in all batch samples
+ return self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionTransformer with OWLVIT->OWLV2,OwlViT->Owlv2
+class Owlv2VisionTransformer(nn.Module):
+ def __init__(self, config: Owlv2VisionConfig):
+ super().__init__()
+ self.config = config
+
+ self.embeddings = Owlv2VisionEmbeddings(config)
+ self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.encoder = Owlv2Encoder(config)
+ self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ @add_start_docstrings_to_model_forward(OWLV2_VISION_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Owlv2VisionConfig)
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: Optional[bool] = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ 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
+
+ # Cast the input to the expected `dtype`
+ expected_input_dtype = self.embeddings.patch_embedding.weight.dtype
+ pixel_values = pixel_values.to(expected_input_dtype)
+
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+ hidden_states = self.pre_layernorm(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]
+ pooled_output = last_hidden_state[:, 0, :]
+
+ pooled_output = self.post_layernorm(pooled_output)
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionModel with OWLVIT->OWLV2,OwlViT->Owlv2,google/owlvit-base-patch32->google/owlv2-base-patch16
+class Owlv2VisionModel(Owlv2PreTrainedModel):
+ config_class = Owlv2VisionConfig
+ main_input_name = "pixel_values"
+
+ def __init__(self, config: Owlv2VisionConfig):
+ super().__init__(config)
+ self.vision_model = Owlv2VisionTransformer(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(OWLV2_VISION_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Owlv2VisionConfig)
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
+ r"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Owlv2VisionModel
+
+ >>> model = Owlv2VisionModel.from_pretrained("google/owlv2-base-patch16")
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16")
+ >>> 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
+ >>> pooled_output = outputs.pooler_output # pooled CLS states
+ ```"""
+ return self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=return_dict,
+ )
+
+
+@add_start_docstrings(OWLV2_START_DOCSTRING)
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTModel with google/owlvit-base-patch32->google/owlv2-base-patch16-ensemble, OWLVIT->OWLV2,OwlViT->Owlv2,owlvit->owlv2,OWL-ViT->OWLv2
+class Owlv2Model(Owlv2PreTrainedModel):
+ config_class = Owlv2Config
+
+ def __init__(self, config: Owlv2Config):
+ super().__init__(config)
+
+ if not isinstance(config.text_config, Owlv2TextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type Owlv2TextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ if not isinstance(config.vision_config, Owlv2VisionConfig):
+ raise TypeError(
+ "config.vision_config is expected to be of type Owlv2VisionConfig but is of type"
+ f" {type(config.vision_config)}."
+ )
+
+ text_config = config.text_config
+ vision_config = config.vision_config
+
+ self.projection_dim = config.projection_dim
+ self.text_embed_dim = text_config.hidden_size
+ self.vision_embed_dim = vision_config.hidden_size
+
+ self.text_model = Owlv2TextTransformer(text_config)
+ self.vision_model = Owlv2VisionTransformer(vision_config)
+
+ self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
+ self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
+ self.logit_scale = nn.Parameter(torch.tensor(config.logit_scale_init_value))
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(OWLV2_TEXT_INPUTS_DOCSTRING)
+ def get_text_features(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_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:
+ text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
+ applying the projection layer to the pooled output of [`Owlv2TextModel`].
+
+ Examples:
+ ```python
+ >>> from transformers import AutoProcessor, Owlv2Model
+
+ >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> inputs = processor(
+ ... text=[["a photo of a cat", "a photo of a dog"], ["photo of a astranaut"]], return_tensors="pt"
+ ... )
+ >>> text_features = model.get_text_features(**inputs)
+ ```"""
+ # Use OWLv2 model's config for some fields (if specified) instead of those of vision & text components.
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # Get embeddings for all text queries in all batch samples
+ text_output = self.text_model(input_ids=input_ids, attention_mask=attention_mask, return_dict=return_dict)
+ pooled_output = text_output[1]
+ text_features = self.text_projection(pooled_output)
+
+ return text_features
+
+ @add_start_docstrings_to_model_forward(OWLV2_VISION_INPUTS_DOCSTRING)
+ def get_image_features(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ 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 [`Owlv2VisionModel`].
+
+ Examples:
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Owlv2Model
+
+ >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> 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)
+ ```"""
+ # Use OWLv2 model's config for some fields (if specified) instead of those of vision & text components.
+ 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
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=return_dict,
+ )
+
+ pooled_output = vision_outputs[1]
+ image_features = self.visual_projection(pooled_output)
+
+ return image_features
+
+ @add_start_docstrings_to_model_forward(OWLV2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Owlv2Output, config_class=Owlv2Config)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ return_loss: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ return_base_image_embeds: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, Owlv2Output]:
+ r"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Owlv2Model
+
+ >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> 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", "a photo of a dog"]], images=image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+ ```"""
+ # Use OWLv2 model's config for some fields (if specified) instead of those of vision & text components.
+ 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
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=return_dict,
+ )
+
+ # Get embeddings for all text queries in all batch samples
+ text_outputs = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ text_embeds = text_outputs[1]
+ text_embeds = self.text_projection(text_embeds)
+ image_embeds = vision_outputs[1]
+ image_embeds = self.visual_projection(image_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / torch.linalg.norm(image_embeds, ord=2, dim=-1, keepdim=True)
+ text_embeds_norm = text_embeds / torch.linalg.norm(text_embeds, ord=2, dim=-1, keepdim=True)
+
+ # cosine similarity as logits and set it on the correct device
+ logit_scale = self.logit_scale.exp().to(image_embeds.device)
+
+ logits_per_text = torch.matmul(text_embeds_norm, image_embeds.t()) * logit_scale
+ logits_per_image = logits_per_text.t()
+
+ loss = None
+ if return_loss:
+ loss = owlv2_loss(logits_per_text)
+
+ text_embeds = text_embeds_norm
+
+ if not return_dict:
+ output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+ return ((loss,) + output) if loss is not None else output
+
+ return Owlv2Output(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTBoxPredictionHead with OwlViT->Owlv2
+class Owlv2BoxPredictionHead(nn.Module):
+ def __init__(self, config: Owlv2Config, out_dim: int = 4):
+ super().__init__()
+
+ width = config.vision_config.hidden_size
+ self.dense0 = nn.Linear(width, width)
+ self.dense1 = nn.Linear(width, width)
+ self.gelu = nn.GELU()
+ self.dense2 = nn.Linear(width, out_dim)
+
+ def forward(self, image_features: torch.Tensor) -> torch.FloatTensor:
+ output = self.dense0(image_features)
+ output = self.gelu(output)
+ output = self.dense1(output)
+ output = self.gelu(output)
+ output = self.dense2(output)
+ return output
+
+
+# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTClassPredictionHead with OwlViT->Owlv2
+class Owlv2ClassPredictionHead(nn.Module):
+ def __init__(self, config: Owlv2Config):
+ super().__init__()
+
+ out_dim = config.text_config.hidden_size
+ self.query_dim = config.vision_config.hidden_size
+
+ self.dense0 = nn.Linear(self.query_dim, out_dim)
+ self.logit_shift = nn.Linear(self.query_dim, 1)
+ self.logit_scale = nn.Linear(self.query_dim, 1)
+ self.elu = nn.ELU()
+
+ def forward(
+ self,
+ image_embeds: torch.FloatTensor,
+ query_embeds: Optional[torch.FloatTensor],
+ query_mask: Optional[torch.Tensor],
+ ) -> Tuple[torch.FloatTensor]:
+ image_class_embeds = self.dense0(image_embeds)
+ if query_embeds is None:
+ device = image_class_embeds.device
+ batch_size, num_patches = image_class_embeds.shape[:2]
+ pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
+ return (pred_logits, image_class_embeds)
+
+ # Normalize image and text features
+ image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)
+ query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)
+
+ # Get class predictions
+ pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)
+
+ # Apply a learnable shift and scale to logits
+ logit_shift = self.logit_shift(image_embeds)
+ logit_scale = self.logit_scale(image_embeds)
+ logit_scale = self.elu(logit_scale) + 1
+ pred_logits = (pred_logits + logit_shift) * logit_scale
+
+ if query_mask is not None:
+ if query_mask.ndim > 1:
+ query_mask = torch.unsqueeze(query_mask, dim=-2)
+
+ pred_logits = torch.where(query_mask == 0, torch.finfo(pred_logits.dtype).min, pred_logits)
+ pred_logits = pred_logits.to(torch.float32)
+
+ return (pred_logits, image_class_embeds)
+
+
+class Owlv2ForObjectDetection(Owlv2PreTrainedModel):
+ config_class = Owlv2Config
+
+ def __init__(self, config: Owlv2Config):
+ super().__init__(config)
+
+ self.owlv2 = Owlv2Model(config)
+ self.class_head = Owlv2ClassPredictionHead(config)
+ self.box_head = Owlv2BoxPredictionHead(config)
+ self.objectness_head = Owlv2BoxPredictionHead(config, out_dim=1)
+
+ self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps)
+ self.sigmoid = nn.Sigmoid()
+ self.config = config
+ self.num_patches_height = self.config.vision_config.image_size // self.config.vision_config.patch_size
+ self.num_patches_width = self.config.vision_config.image_size // self.config.vision_config.patch_size
+ self.box_bias = self.compute_box_bias(self.num_patches_height, self.num_patches_width)
+
+ @staticmethod
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.normalize_grid_corner_coordinates
+ def normalize_grid_corner_coordinates(num_patches_height: int, num_patches_width: int) -> torch.Tensor:
+ # Create grid coordinates using torch
+ x_coordinates = torch.arange(1, num_patches_width + 1, dtype=torch.float32)
+ y_coordinates = torch.arange(1, num_patches_height + 1, dtype=torch.float32)
+ xx, yy = torch.meshgrid(x_coordinates, y_coordinates, indexing="xy")
+
+ # Stack the coordinates and divide by their respective patch counts
+ box_coordinates = torch.stack((xx, yy), dim=-1)
+ box_coordinates[..., 0] /= num_patches_width
+ box_coordinates[..., 1] /= num_patches_height
+
+ # Flatten (h, w, 2) -> (h*w, 2)
+ box_coordinates = box_coordinates.view(-1, 2)
+
+ return box_coordinates
+
+ def objectness_predictor(self, image_features: torch.FloatTensor) -> torch.FloatTensor:
+ """Predicts the probability that each image feature token is an object.
+
+ Args:
+ image_features (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_dim)`)):
+ Features extracted from the image.
+ Returns:
+ Objectness scores.
+ """
+ image_features = image_features.detach()
+ objectness_logits = self.objectness_head(image_features)
+ objectness_logits = objectness_logits[..., 0]
+ return objectness_logits
+
+ @lru_cache(maxsize=2)
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.compute_box_bias
+ def compute_box_bias(
+ self, num_patches_height: int, num_patches_width: int, feature_map: Optional[torch.FloatTensor] = None
+ ) -> torch.Tensor:
+ if feature_map is not None:
+ raise ValueError("feature_map has been deprecated as an input. Please pass in num_patches instead")
+ # The box center is biased to its position on the feature grid
+ box_coordinates = self.normalize_grid_corner_coordinates(num_patches_height, num_patches_width)
+ box_coordinates = torch.clip(box_coordinates, 0.0, 1.0)
+
+ # Unnormalize xy
+ box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4)
+
+ # The box size is biased to the patch size
+ box_size = torch.full_like(box_coord_bias, 1.0)
+ box_size[..., 0] /= num_patches_width
+ box_size[..., 1] /= num_patches_height
+ box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4)
+
+ # Compute box bias
+ box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1)
+ return box_bias
+
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.box_predictor
+ def box_predictor(
+ self,
+ image_feats: torch.FloatTensor,
+ feature_map: torch.FloatTensor,
+ interpolate_pos_encoding: bool = False,
+ ) -> torch.FloatTensor:
+ """
+ Args:
+ image_feats:
+ Features extracted from the image, returned by the `image_text_embedder` method.
+ feature_map:
+ A spatial re-arrangement of image_features, also returned by the `image_text_embedder` method.
+ interpolate_pos_encoding:
+ Whether to interpolate the pre-trained position encodings.
+ Returns:
+ pred_boxes:
+ List of predicted boxes (cxcywh normalized to 0, 1) nested within a dictionary.
+ """
+ # Bounding box detection head [batch_size, num_boxes, 4].
+ pred_boxes = self.box_head(image_feats)
+
+ # Compute the location of each token on the grid and use it to compute a bias for the bbox prediction
+ if interpolate_pos_encoding:
+ _, num_patches_height, num_patches_width, _ = feature_map.shape
+ box_bias = self.compute_box_bias(num_patches_height, num_patches_width)
+ else:
+ box_bias = self.box_bias
+
+ box_bias = box_bias.to(feature_map.device)
+ pred_boxes += box_bias
+ pred_boxes = self.sigmoid(pred_boxes)
+ return pred_boxes
+
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.class_predictor
+ def class_predictor(
+ self,
+ image_feats: torch.FloatTensor,
+ query_embeds: Optional[torch.FloatTensor] = None,
+ query_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.FloatTensor]:
+ """
+ Args:
+ image_feats:
+ Features extracted from the `image_text_embedder`.
+ query_embeds:
+ Text query embeddings.
+ query_mask:
+ Must be provided with query_embeddings. A mask indicating which query embeddings are valid.
+ """
+ (pred_logits, image_class_embeds) = self.class_head(image_feats, query_embeds, query_mask)
+
+ return (pred_logits, image_class_embeds)
+
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.image_text_embedder with owlvit->owlv2
+ def image_text_embedder(
+ self,
+ input_ids: torch.Tensor,
+ pixel_values: torch.FloatTensor,
+ attention_mask: torch.Tensor,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> Tuple[torch.FloatTensor]:
+ # Encode text and image
+ outputs = self.owlv2(
+ pixel_values=pixel_values,
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ )
+
+ if interpolate_pos_encoding:
+ _, _, height, width = pixel_values.shape
+ num_patches_height = height // self.config.vision_config.patch_size
+ num_patches_width = width // self.config.vision_config.patch_size
+ else:
+ num_patches_height = self.num_patches_height
+ num_patches_width = self.num_patches_width
+
+ # Get image embeddings
+ last_hidden_state = outputs.vision_model_output[0]
+ image_embeds = self.owlv2.vision_model.post_layernorm(last_hidden_state)
+
+ # Resize class token
+ class_token_out = torch.broadcast_to(image_embeds[:, :1, :], image_embeds[:, :-1].shape)
+
+ # Merge image embedding with class tokens
+ image_embeds = image_embeds[:, 1:, :] * class_token_out
+ image_embeds = self.layer_norm(image_embeds)
+
+ # Resize to [batch_size, num_patches_height, num_patches_width, hidden_size]
+ new_size = (
+ image_embeds.shape[0],
+ num_patches_height,
+ num_patches_width,
+ image_embeds.shape[-1],
+ )
+ image_embeds = image_embeds.reshape(new_size)
+ text_embeds = outputs[-4]
+
+ return (text_embeds, image_embeds, outputs)
+
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.image_embedder with owlvit->owlv2, OwlViTModel->Owlv2Model
+ def image_embedder(
+ self,
+ pixel_values: torch.FloatTensor,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> Tuple[torch.FloatTensor]:
+ # Get Owlv2Model vision embeddings (same as CLIP)
+ vision_outputs = self.owlv2.vision_model(
+ pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True
+ )
+
+ if interpolate_pos_encoding:
+ _, _, height, width = pixel_values.shape
+ num_patches_height = height // self.config.vision_config.patch_size
+ num_patches_width = width // self.config.vision_config.patch_size
+ else:
+ num_patches_height = self.num_patches_height
+ num_patches_width = self.num_patches_width
+
+ # Apply post_layernorm to last_hidden_state, return non-projected output
+ last_hidden_state = vision_outputs[0]
+ image_embeds = self.owlv2.vision_model.post_layernorm(last_hidden_state)
+
+ # Resize class token
+ class_token_out = torch.broadcast_to(image_embeds[:, :1, :], image_embeds[:, :-1].shape)
+
+ # Merge image embedding with class tokens
+ image_embeds = image_embeds[:, 1:, :] * class_token_out
+ image_embeds = self.layer_norm(image_embeds)
+
+ # Resize to [batch_size, num_patches_height, num_patches_width, hidden_size]
+ new_size = (
+ image_embeds.shape[0],
+ num_patches_height,
+ num_patches_width,
+ image_embeds.shape[-1],
+ )
+ image_embeds = image_embeds.reshape(new_size)
+
+ return (image_embeds, vision_outputs)
+
+ # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTForObjectDetection.embed_image_query
+ def embed_image_query(
+ self,
+ query_image_features: torch.FloatTensor,
+ query_feature_map: torch.FloatTensor,
+ interpolate_pos_encoding: bool = False,
+ ) -> torch.FloatTensor:
+ _, class_embeds = self.class_predictor(query_image_features)
+ pred_boxes = self.box_predictor(query_image_features, query_feature_map, interpolate_pos_encoding)
+ pred_boxes_as_corners = center_to_corners_format(pred_boxes)
+
+ # Loop over query images
+ best_class_embeds = []
+ best_box_indices = []
+ pred_boxes_device = pred_boxes_as_corners.device
+
+ for i in range(query_image_features.shape[0]):
+ each_query_box = torch.tensor([[0, 0, 1, 1]], device=pred_boxes_device)
+ each_query_pred_boxes = pred_boxes_as_corners[i]
+ ious, _ = box_iou(each_query_box, each_query_pred_boxes)
+
+ # If there are no overlapping boxes, fall back to generalized IoU
+ if torch.all(ious[0] == 0.0):
+ ious = generalized_box_iou(each_query_box, each_query_pred_boxes)
+
+ # Use an adaptive threshold to include all boxes within 80% of the best IoU
+ iou_threshold = torch.max(ious) * 0.8
+
+ selected_inds = (ious[0] >= iou_threshold).nonzero()
+ if selected_inds.numel():
+ selected_embeddings = class_embeds[i][selected_inds.squeeze(1)]
+ mean_embeds = torch.mean(class_embeds[i], axis=0)
+ mean_sim = torch.einsum("d,id->i", mean_embeds, selected_embeddings)
+ best_box_ind = selected_inds[torch.argmin(mean_sim)]
+ best_class_embeds.append(class_embeds[i][best_box_ind])
+ best_box_indices.append(best_box_ind)
+
+ if best_class_embeds:
+ query_embeds = torch.stack(best_class_embeds)
+ box_indices = torch.stack(best_box_indices)
+ else:
+ query_embeds, box_indices = None, None
+
+ return query_embeds, box_indices, pred_boxes
+
+ @add_start_docstrings_to_model_forward(OWLV2_IMAGE_GUIDED_OBJECT_DETECTION_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Owlv2ImageGuidedObjectDetectionOutput, config_class=Owlv2Config)
+ def image_guided_detection(
+ self,
+ pixel_values: torch.FloatTensor,
+ query_pixel_values: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Owlv2ImageGuidedObjectDetectionOutput:
+ r"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> import requests
+ >>> from PIL import Image
+ >>> import torch
+ >>> from transformers import AutoProcessor, Owlv2ForObjectDetection
+
+ >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+ >>> query_url = "http://images.cocodataset.org/val2017/000000001675.jpg"
+ >>> query_image = Image.open(requests.get(query_url, stream=True).raw)
+ >>> inputs = processor(images=image, query_images=query_image, return_tensors="pt")
+
+ >>> # forward pass
+ >>> with torch.no_grad():
+ ... outputs = model.image_guided_detection(**inputs)
+
+ >>> target_sizes = torch.Tensor([image.size[::-1]])
+
+ >>> # Convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+ >>> results = processor.post_process_image_guided_detection(
+ ... outputs=outputs, threshold=0.9, nms_threshold=0.3, target_sizes=target_sizes
+ ... )
+ >>> i = 0 # Retrieve predictions for the first image
+ >>> boxes, scores = results[i]["boxes"], results[i]["scores"]
+ >>> for box, score in zip(boxes, scores):
+ ... box = [round(i, 2) for i in box.tolist()]
+ ... print(f"Detected similar object with confidence {round(score.item(), 3)} at location {box}")
+ Detected similar object with confidence 0.938 at location [327.31, 54.94, 547.39, 268.06]
+ Detected similar object with confidence 0.959 at location [5.78, 360.65, 619.12, 366.39]
+ Detected similar object with confidence 0.902 at location [2.85, 360.01, 627.63, 380.8]
+ Detected similar object with confidence 0.985 at location [176.98, -29.45, 672.69, 182.83]
+ Detected similar object with confidence 1.0 at location [6.53, 14.35, 624.87, 470.82]
+ Detected similar object with confidence 0.998 at location [579.98, 29.14, 615.49, 489.05]
+ Detected similar object with confidence 0.985 at location [206.15, 10.53, 247.74, 466.01]
+ Detected similar object with confidence 0.947 at location [18.62, 429.72, 646.5, 457.72]
+ Detected similar object with confidence 0.996 at location [523.88, 20.69, 586.84, 483.18]
+ Detected similar object with confidence 0.998 at location [3.39, 360.59, 617.29, 499.21]
+ Detected similar object with confidence 0.969 at location [4.47, 449.05, 614.5, 474.76]
+ Detected similar object with confidence 0.966 at location [31.44, 463.65, 654.66, 471.07]
+ Detected similar object with confidence 0.924 at location [30.93, 468.07, 635.35, 475.39]
+ ```"""
+ 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
+
+ # Compute feature maps for the input and query images
+ query_feature_map = self.image_embedder(
+ pixel_values=query_pixel_values, interpolate_pos_encoding=interpolate_pos_encoding
+ )[0]
+ feature_map, vision_outputs = self.image_embedder(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ batch_size, num_patches_height, num_patches_width, hidden_dim = feature_map.shape
+ image_feats = torch.reshape(feature_map, (batch_size, num_patches_height * num_patches_width, hidden_dim))
+
+ batch_size, num_patches_height, num_patches_width, hidden_dim = query_feature_map.shape
+ query_image_feats = torch.reshape(
+ query_feature_map, (batch_size, num_patches_height * num_patches_width, hidden_dim)
+ )
+ # Get top class embedding and best box index for each query image in batch
+ query_embeds, best_box_indices, query_pred_boxes = self.embed_image_query(
+ query_image_feats, query_feature_map, interpolate_pos_encoding
+ )
+
+ # Predict object classes [batch_size, num_patches, num_queries+1]
+ (pred_logits, class_embeds) = self.class_predictor(image_feats=image_feats, query_embeds=query_embeds)
+
+ # Predict object boxes
+ target_pred_boxes = self.box_predictor(image_feats, feature_map, interpolate_pos_encoding)
+
+ if not return_dict:
+ output = (
+ feature_map,
+ query_feature_map,
+ target_pred_boxes,
+ query_pred_boxes,
+ pred_logits,
+ class_embeds,
+ vision_outputs.to_tuple(),
+ )
+ output = tuple(x for x in output if x is not None)
+ return output
+
+ return Owlv2ImageGuidedObjectDetectionOutput(
+ image_embeds=feature_map,
+ query_image_embeds=query_feature_map,
+ target_pred_boxes=target_pred_boxes,
+ query_pred_boxes=query_pred_boxes,
+ logits=pred_logits,
+ class_embeds=class_embeds,
+ text_model_output=None,
+ vision_model_output=vision_outputs,
+ )
+
+ @add_start_docstrings_to_model_forward(OWLV2_OBJECT_DETECTION_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Owlv2ObjectDetectionOutput, config_class=Owlv2Config)
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ pixel_values: torch.FloatTensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Owlv2ObjectDetectionOutput:
+ r"""
+ Returns:
+
+ Examples:
+ ```python
+ >>> import requests
+ >>> from PIL import Image
+ >>> import torch
+
+ >>> from transformers import Owlv2Processor, Owlv2ForObjectDetection
+
+ >>> processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+ >>> model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+ >>> text_labels = [["a photo of a cat", "a photo of a dog"]]
+ >>> inputs = processor(text=text_labels, images=image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> # Target image sizes (height, width) to rescale box predictions [batch_size, 2]
+ >>> target_sizes = torch.tensor([(image.height, image.width)])
+ >>> # Convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+ >>> results = processor.post_process_grounded_object_detection(
+ ... outputs=outputs, target_sizes=target_sizes, threshold=0.1, text_labels=text_labels
+ ... )
+ >>> # Retrieve predictions for the first image for the corresponding text queries
+ >>> result = results[0]
+ >>> boxes, scores, text_labels = result["boxes"], result["scores"], result["text_labels"]
+ >>> for box, score, text_label in zip(boxes, scores, text_labels):
+ ... box = [round(i, 2) for i in box.tolist()]
+ ... print(f"Detected {text_label} with confidence {round(score.item(), 3)} at location {box}")
+ Detected a photo of a cat with confidence 0.614 at location [341.67, 23.39, 642.32, 371.35]
+ Detected a photo of a cat with confidence 0.665 at location [6.75, 51.96, 326.62, 473.13]
+ ```"""
+ 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
+
+ # Embed images and text queries
+ query_embeds, feature_map, outputs = self.image_text_embedder(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ )
+
+ # Text and vision model outputs
+ text_outputs = outputs.text_model_output
+ vision_outputs = outputs.vision_model_output
+
+ batch_size, num_patches_height, num_patches_width, hidden_dim = feature_map.shape
+ image_feats = torch.reshape(feature_map, (batch_size, num_patches_height * num_patches_width, hidden_dim))
+
+ # Reshape from [batch_size * max_text_queries, hidden_dim] -> [batch_size, max_text_queries, hidden_dim]
+ max_text_queries = input_ids.shape[0] // batch_size
+ query_embeds = query_embeds.reshape(batch_size, max_text_queries, query_embeds.shape[-1])
+
+ # If first token is 0, then this is a padded query [batch_size, num_queries].
+ input_ids = input_ids.reshape(batch_size, max_text_queries, input_ids.shape[-1])
+ query_mask = input_ids[..., 0] > 0
+
+ # Predict object classes [batch_size, num_patches, num_queries+1]
+ (pred_logits, class_embeds) = self.class_predictor(image_feats, query_embeds, query_mask)
+
+ # Predict objectness
+ objectness_logits = self.objectness_predictor(image_feats)
+
+ # Predict object boxes
+ pred_boxes = self.box_predictor(image_feats, feature_map, interpolate_pos_encoding)
+
+ if not return_dict:
+ output = (
+ pred_logits,
+ objectness_logits,
+ pred_boxes,
+ query_embeds,
+ feature_map,
+ class_embeds,
+ text_outputs.to_tuple(),
+ vision_outputs.to_tuple(),
+ )
+ output = tuple(x for x in output if x is not None)
+ return output
+
+ return Owlv2ObjectDetectionOutput(
+ image_embeds=feature_map,
+ text_embeds=query_embeds,
+ pred_boxes=pred_boxes,
+ logits=pred_logits,
+ objectness_logits=objectness_logits,
+ class_embeds=class_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+__all__ = ["Owlv2Model", "Owlv2PreTrainedModel", "Owlv2TextModel", "Owlv2VisionModel", "Owlv2ForObjectDetection"]
diff --git a/lib/python3.12/site-packages/transformers/models/owlv2/processing_owlv2.py b/lib/python3.12/site-packages/transformers/models/owlv2/processing_owlv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..4996cae7ab9694e2114cdc046ca95837a870158f
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/owlv2/processing_owlv2.py
@@ -0,0 +1,320 @@
+# 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 OWLv2
+"""
+
+import warnings
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ...image_processing_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import (
+ ImagesKwargs,
+ ProcessingKwargs,
+ ProcessorMixin,
+ Unpack,
+ _validate_images_text_input_order,
+)
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import TensorType, is_flax_available, is_tf_available, is_torch_available
+
+
+if TYPE_CHECKING:
+ from .modeling_owlv2 import Owlv2ImageGuidedObjectDetectionOutput, Owlv2ObjectDetectionOutput
+
+
+class Owlv2ImagesKwargs(ImagesKwargs, total=False):
+ query_images: Optional[ImageInput]
+
+
+class Owlv2ProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: Owlv2ImagesKwargs
+ _defaults = {
+ "text_kwargs": {
+ "padding": "max_length",
+ },
+ "images_kwargs": {},
+ "common_kwargs": {
+ "return_tensors": "np",
+ },
+ }
+
+
+class Owlv2Processor(ProcessorMixin):
+ r"""
+ Constructs an Owlv2 processor which wraps [`Owlv2ImageProcessor`] and [`CLIPTokenizer`]/[`CLIPTokenizerFast`] into
+ a single processor that interits both the image processor and tokenizer functionalities. See the
+ [`~OwlViTProcessor.__call__`] and [`~OwlViTProcessor.decode`] for more information.
+
+ Args:
+ image_processor ([`Owlv2ImageProcessor`]):
+ The image processor is a required input.
+ tokenizer ([`CLIPTokenizer`, `CLIPTokenizerFast`]):
+ The tokenizer is a required input.
+ """
+
+ attributes = ["image_processor", "tokenizer"]
+ image_processor_class = "Owlv2ImageProcessor"
+ tokenizer_class = ("CLIPTokenizer", "CLIPTokenizerFast")
+ # For backward compatibility. See transformers.processing_utils.ProcessorMixin.prepare_and_validate_optional_call_args for more details.
+ optional_call_args = ["query_images"]
+
+ def __init__(self, image_processor, tokenizer, **kwargs):
+ super().__init__(image_processor, tokenizer)
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.__call__ with OwlViT->Owlv2
+ def __call__(
+ self,
+ images: Optional[ImageInput] = None,
+ text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+ # The following is to capture `query_images` argument that may be passed as a positional argument.
+ # See transformers.processing_utils.ProcessorMixin.prepare_and_validate_optional_call_args for more details,
+ # or this conversation for more context: https://github.com/huggingface/transformers/pull/32544#discussion_r1720208116
+ # This behavior is only needed for backward compatibility and will be removed in future versions.
+ #
+ *args,
+ audio=None,
+ videos=None,
+ **kwargs: Unpack[Owlv2ProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several text(s) and image(s). This method forwards the `text` and
+ `kwargs` arguments to CLIPTokenizerFast's [`~CLIPTokenizerFast.__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 docstring
+ of the above two methods for more information.
+
+ Args:
+ 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.
+ 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).
+ query_images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The query image to be prepared, one query image is expected per target image to be queried. Each image
+ can be a PIL image, NumPy array or PyTorch tensor. In case of a NumPy array/PyTorch tensor, each image
+ should be of shape (C, H, W), where C is a number of channels, H and W are image height and width.
+ 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:
+ [`BatchFeature`]: A [`BatchFeature`] 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`.
+ - **query_pixel_values** -- Pixel values of the query images to be fed to a model. Returned when `query_images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ Owlv2ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ **self.prepare_and_validate_optional_call_args(*args),
+ )
+ query_images = output_kwargs["images_kwargs"].pop("query_images", None)
+ return_tensors = output_kwargs["common_kwargs"]["return_tensors"]
+
+ if text is None and query_images is None and images is None:
+ raise ValueError(
+ "You have to specify at least one text or query image or image. All three cannot be none."
+ )
+ # check if images and text inputs are reversed for BC
+ images, text = _validate_images_text_input_order(images, text)
+
+ data = {}
+ if text is not None:
+ if isinstance(text, str) or (isinstance(text, List) and not isinstance(text[0], List)):
+ encodings = [self.tokenizer(text, **output_kwargs["text_kwargs"])]
+
+ elif isinstance(text, List) and isinstance(text[0], List):
+ encodings = []
+
+ # Maximum number of queries across batch
+ max_num_queries = max([len(text_single) for text_single in text])
+
+ # Pad all batch samples to max number of text queries
+ for text_single in text:
+ if len(text_single) != max_num_queries:
+ text_single = text_single + [" "] * (max_num_queries - len(text_single))
+
+ encoding = self.tokenizer(text_single, **output_kwargs["text_kwargs"])
+ encodings.append(encoding)
+ else:
+ raise TypeError("Input text should be a string, a list of strings or a nested list of strings")
+
+ if return_tensors == "np":
+ input_ids = np.concatenate([encoding["input_ids"] for encoding in encodings], axis=0)
+ attention_mask = np.concatenate([encoding["attention_mask"] for encoding in encodings], axis=0)
+
+ elif return_tensors == "jax" and is_flax_available():
+ import jax.numpy as jnp
+
+ input_ids = jnp.concatenate([encoding["input_ids"] for encoding in encodings], axis=0)
+ attention_mask = jnp.concatenate([encoding["attention_mask"] for encoding in encodings], axis=0)
+
+ elif return_tensors == "pt" and is_torch_available():
+ import torch
+
+ input_ids = torch.cat([encoding["input_ids"] for encoding in encodings], dim=0)
+ attention_mask = torch.cat([encoding["attention_mask"] for encoding in encodings], dim=0)
+
+ elif return_tensors == "tf" and is_tf_available():
+ import tensorflow as tf
+
+ input_ids = tf.stack([encoding["input_ids"] for encoding in encodings], axis=0)
+ attention_mask = tf.stack([encoding["attention_mask"] for encoding in encodings], axis=0)
+
+ else:
+ raise ValueError("Target return tensor type could not be returned")
+
+ data["input_ids"] = input_ids
+ data["attention_mask"] = attention_mask
+
+ if query_images is not None:
+ query_pixel_values = self.image_processor(query_images, **output_kwargs["images_kwargs"]).pixel_values
+ # Query images always override the text prompt
+ data = {"query_pixel_values": query_pixel_values}
+
+ if images is not None:
+ image_features = self.image_processor(images, **output_kwargs["images_kwargs"])
+ data["pixel_values"] = image_features.pixel_values
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.post_process_object_detection with OwlViT->Owlv2
+ def post_process_object_detection(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to [`Owlv2ImageProcessor.post_process_object_detection`]. Please refer
+ to the docstring of this method for more information.
+ """
+ warnings.warn(
+ "`post_process_object_detection` method is deprecated for OwlVitProcessor and will be removed in v5. "
+ "Use `post_process_grounded_object_detection` instead.",
+ FutureWarning,
+ )
+ return self.image_processor.post_process_object_detection(*args, **kwargs)
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.post_process_grounded_object_detection with OwlViT->Owlv2
+ def post_process_grounded_object_detection(
+ self,
+ outputs: "Owlv2ObjectDetectionOutput",
+ threshold: float = 0.1,
+ target_sizes: Optional[Union[TensorType, List[Tuple]]] = None,
+ text_labels: Optional[List[List[str]]] = None,
+ ):
+ """
+ Converts the raw output of [`Owlv2ForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
+ bottom_right_x, bottom_right_y) format.
+
+ Args:
+ outputs ([`Owlv2ObjectDetectionOutput`]):
+ Raw outputs of the model.
+ threshold (`float`, *optional*, defaults to 0.1):
+ 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.
+ text_labels (`List[List[str]]`, *optional*):
+ List of lists of text labels for each image in the batch. If unset, "text_labels" in output will be
+ set to `None`.
+
+ Returns:
+ `List[Dict]`: A list of dictionaries, each dictionary containing the following keys:
+ - "scores": The confidence scores for each predicted box on the image.
+ - "labels": Indexes of the classes predicted by the model on the image.
+ - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.
+ - "text_labels": The text labels for each predicted bounding box on the image.
+ """
+ output = self.image_processor.post_process_object_detection(
+ outputs=outputs, threshold=threshold, target_sizes=target_sizes
+ )
+
+ if text_labels is not None and len(text_labels) != len(output):
+ raise ValueError("Make sure that you pass in as many lists of text labels as images")
+
+ # adding text labels to the output
+ if text_labels is not None:
+ for image_output, image_text_labels in zip(output, text_labels):
+ object_text_labels = [image_text_labels[i] for i in image_output["labels"]]
+ image_output["text_labels"] = object_text_labels
+ else:
+ for image_output in output:
+ image_output["text_labels"] = None
+
+ return output
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.post_process_image_guided_detection with OwlViT->Owlv2
+ def post_process_image_guided_detection(
+ self,
+ outputs: "Owlv2ImageGuidedObjectDetectionOutput",
+ threshold: float = 0.0,
+ nms_threshold: float = 0.3,
+ target_sizes: Optional[Union[TensorType, List[Tuple]]] = None,
+ ):
+ """
+ Converts the output of [`Owlv2ForObjectDetection.image_guided_detection`] into the format expected by the COCO
+ api.
+
+ Args:
+ outputs ([`Owlv2ImageGuidedObjectDetectionOutput`]):
+ Raw outputs of the model.
+ threshold (`float`, *optional*, defaults to 0.0):
+ Minimum confidence threshold to use to filter out predicted boxes.
+ nms_threshold (`float`, *optional*, defaults to 0.3):
+ IoU threshold for non-maximum suppression of overlapping boxes.
+ 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, predicted normalized bounding boxes are rescaled to the target sizes. If left to
+ None, predictions will not be unnormalized.
+
+ Returns:
+ `List[Dict]`: A list of dictionaries, each dictionary containing the following keys:
+ - "scores": The confidence scores for each predicted box on the image.
+ - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.
+ - "labels": Set to `None`.
+ """
+ return self.image_processor.post_process_image_guided_detection(
+ outputs=outputs, threshold=threshold, nms_threshold=nms_threshold, target_sizes=target_sizes
+ )
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.batch_decode
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+ refer to the docstring of this method for more information.
+ """
+ return self.tokenizer.batch_decode(*args, **kwargs)
+
+ # Copied from transformers.models.owlvit.processing_owlvit.OwlViTProcessor.decode
+ def decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+ the docstring of this method for more information.
+ """
+ return self.tokenizer.decode(*args, **kwargs)
+
+
+__all__ = ["Owlv2Processor"]
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/__init__.py b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..285c1970308a47827806fca349d130703f40a2c8
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_patchtsmixer import *
+ from .modeling_patchtsmixer import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e606e23c2df5e18e9903839ad3d15fdef5b4a410
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/configuration_patchtsmixer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/configuration_patchtsmixer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..71c938f77d13262a6def6358f45d1b5e432d2e92
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/configuration_patchtsmixer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/modeling_patchtsmixer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/modeling_patchtsmixer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5c2e0c39a76ac122d0a65267067da9d9f97d7877
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/patchtsmixer/__pycache__/modeling_patchtsmixer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/configuration_patchtsmixer.py b/lib/python3.12/site-packages/transformers/models/patchtsmixer/configuration_patchtsmixer.py
new file mode 100644
index 0000000000000000000000000000000000000000..83f374651a7cfc998349a734d8350fc399ea7f07
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/patchtsmixer/configuration_patchtsmixer.py
@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2023 IBM 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.
+"""PatchTSMixer model configuration"""
+
+from typing import List, Optional, Union
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class PatchTSMixerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`PatchTSMixerModel`]. It is used to instantiate a
+ PatchTSMixer 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 PatchTSMixer
+ [ibm/patchtsmixer-etth1-pretrain](https://huggingface.co/ibm/patchtsmixer-etth1-pretrain) architecture.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ context_length (`int`, *optional*, defaults to 32):
+ The context/history length for the input sequence.
+ patch_length (`int`, *optional*, defaults to 8):
+ The patch length for the input sequence.
+ num_input_channels (`int`, *optional*, defaults to 1):
+ Number of input variates. For Univariate, set it to 1.
+ patch_stride (`int`, *optional*, defaults to 8):
+ Determines the overlap between two consecutive patches. Set it to patch_length (or greater), if we want
+ non-overlapping patches.
+ num_parallel_samples (`int`, *optional*, defaults to 100):
+ The number of samples to generate in parallel for probabilistic forecast.
+ d_model (`int`, *optional*, defaults to 8):
+ Hidden dimension of the model. Recommended to set it as a multiple of patch_length (i.e. 2-5X of
+ patch_length). Larger value indicates more complex model.
+ expansion_factor (`int`, *optional*, defaults to 2):
+ Expansion factor to use inside MLP. Recommended range is 2-5. Larger value indicates more complex model.
+ num_layers (`int`, *optional*, defaults to 3):
+ Number of layers to use. Recommended range is 3-15. Larger value indicates more complex model.
+ dropout (`float`, *optional*, defaults to 0.2):
+ The dropout probability the `PatchTSMixer` backbone. Recommended range is 0.2-0.7
+ mode (`str`, *optional*, defaults to `"common_channel"`):
+ Mixer Mode. Determines how to process the channels. Allowed values: "common_channel", "mix_channel". In
+ "common_channel" mode, we follow Channel-independent modelling with no explicit channel-mixing. Channel
+ mixing happens in an implicit manner via shared weights across channels. (preferred first approach) In
+ "mix_channel" mode, we follow explicit channel-mixing in addition to patch and feature mixer. (preferred
+ approach when channel correlations are very important to model)
+ gated_attn (`bool`, *optional*, defaults to `True`):
+ Enable Gated Attention.
+ norm_mlp (`str`, *optional*, defaults to `"LayerNorm"`):
+ Normalization layer (BatchNorm or LayerNorm).
+ self_attn (`bool`, *optional*, defaults to `False`):
+ Enable Tiny self attention across patches. This can be enabled when the output of Vanilla PatchTSMixer with
+ gated attention is not satisfactory. Enabling this leads to explicit pair-wise attention and modelling
+ across patches.
+ self_attn_heads (`int`, *optional*, defaults to 1):
+ Number of self-attention heads. Works only when `self_attn` is set to `True`.
+ use_positional_encoding (`bool`, *optional*, defaults to `False`):
+ Enable the use of positional embedding for the tiny self-attention layers. Works only when `self_attn` is
+ set to `True`.
+ positional_encoding_type (`str`, *optional*, defaults to `"sincos"`):
+ Positional encodings. Options `"random"` and `"sincos"` are supported. Works only when
+ `use_positional_encoding` is set to `True`
+ scaling (`string` or `bool`, *optional*, defaults to `"std"`):
+ Whether to scale the input targets via "mean" scaler, "std" scaler or no scaler if `None`. If `True`, the
+ scaler is set to "mean".
+ loss (`string`, *optional*, defaults to `"mse"`):
+ The loss function for the model corresponding to the `distribution_output` head. For parametric
+ distributions it is the negative log likelihood ("nll") and for point estimates it is the mean squared
+ error "mse".
+ init_std (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated normal weight initialization distribution.
+ post_init (`bool`, *optional*, defaults to `False`):
+ Whether to use custom weight initialization from `transformers` library, or the default initialization in
+ `PyTorch`. Setting it to `False` performs `PyTorch` weight initialization.
+ norm_eps (`float`, *optional*, defaults to 1e-05):
+ A value added to the denominator for numerical stability of normalization.
+ mask_type (`str`, *optional*, defaults to `"random"`):
+ Type of masking to use for Masked Pretraining mode. Allowed values are "random", "forecast". In Random
+ masking, points are masked randomly. In Forecast masking, points are masked towards the end.
+ random_mask_ratio (`float`, *optional*, defaults to 0.5):
+ Masking ratio to use when `mask_type` is `random`. Higher value indicates more masking.
+ num_forecast_mask_patches (`int` or `list`, *optional*, defaults to `[2]`):
+ Number of patches to be masked at the end of each batch sample. If it is an integer, all the samples in the
+ batch will have the same number of masked patches. If it is a list, samples in the batch will be randomly
+ masked by numbers defined in the list. This argument is only used for forecast pretraining.
+ mask_value (`float`, *optional*, defaults to `0.0`):
+ Mask value to use.
+ masked_loss (`bool`, *optional*, defaults to `True`):
+ Whether to compute pretraining loss only at the masked portions, or on the entire output.
+ channel_consistent_masking (`bool`, *optional*, defaults to `True`):
+ When true, masking will be same across all channels of a timeseries. Otherwise, masking positions will vary
+ across channels.
+ unmasked_channel_indices (`list`, *optional*):
+ Channels that are not masked during pretraining.
+ head_dropout (`float`, *optional*, defaults to 0.2):
+ The dropout probability the `PatchTSMixer` head.
+ distribution_output (`string`, *optional*, defaults to `"student_t"`):
+ The distribution emission head for the model when loss is "nll". Could be either "student_t", "normal" or
+ "negative_binomial".
+ prediction_length (`int`, *optional*, defaults to 16):
+ Number of time steps to forecast for a forecasting task. Also known as the Forecast Horizon.
+ prediction_channel_indices (`list`, *optional*):
+ List of channel indices to forecast. If None, forecast all channels. Target data is expected to have all
+ channels and we explicitly filter the channels in prediction and target before loss computation.
+ num_targets (`int`, *optional*, defaults to 3):
+ Number of targets (dimensionality of the regressed variable) for a regression task.
+ output_range (`list`, *optional*):
+ Output range to restrict for the regression task. Defaults to None.
+ head_aggregation (`str`, *optional*, defaults to `"max_pool"`):
+ Aggregation mode to enable for classification or regression task. Allowed values are `None`, "use_last",
+ "max_pool", "avg_pool".
+
+ Example:
+
+ ```python
+ >>> from transformers import PatchTSMixerConfig, PatchTSMixerModel
+
+ >>> # Initializing a default PatchTSMixer configuration
+ >>> configuration = PatchTSMixerConfig()
+
+ >>> # Randomly initializing a model (with random weights) from the configuration
+ >>> model = PatchTSMixerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "patchtsmixer"
+ attribute_map = {
+ "hidden_size": "d_model",
+ "num_hidden_layers": "num_layers",
+ }
+
+ def __init__(
+ self,
+ # Time series specific configuration
+ context_length: int = 32,
+ patch_length: int = 8,
+ num_input_channels: int = 1,
+ patch_stride: int = 8,
+ num_parallel_samples: int = 100,
+ # General model configuration
+ d_model: int = 8,
+ expansion_factor: int = 2,
+ num_layers: int = 3,
+ dropout: float = 0.2,
+ mode: str = "common_channel",
+ gated_attn: bool = True,
+ norm_mlp: str = "LayerNorm",
+ self_attn: bool = False,
+ self_attn_heads: int = 1,
+ use_positional_encoding: bool = False,
+ positional_encoding_type: str = "sincos",
+ scaling: Optional[Union[str, bool]] = "std",
+ loss: str = "mse",
+ init_std: float = 0.02,
+ post_init: bool = False,
+ norm_eps: float = 1e-5,
+ # Pretrain model configuration
+ mask_type: str = "random",
+ random_mask_ratio: float = 0.5,
+ num_forecast_mask_patches: Optional[Union[List[int], int]] = [2],
+ mask_value: int = 0,
+ masked_loss: bool = True,
+ channel_consistent_masking: bool = True,
+ unmasked_channel_indices: Optional[List[int]] = None,
+ # General head configuration
+ head_dropout: float = 0.2,
+ distribution_output: str = "student_t",
+ # Prediction head configuration
+ prediction_length: int = 16,
+ prediction_channel_indices: list = None,
+ # Classification/Regression configuration
+ num_targets: int = 3,
+ output_range: list = None,
+ head_aggregation: str = "max_pool",
+ **kwargs,
+ ):
+ self.num_input_channels = num_input_channels
+ self.context_length = context_length
+ self.patch_length = patch_length
+ self.patch_stride = patch_stride
+ self.d_model = d_model
+ self.expansion_factor = expansion_factor
+ self.num_layers = num_layers
+ self.dropout = dropout
+ self.mode = mode
+ self.gated_attn = gated_attn
+ self.norm_mlp = norm_mlp
+ self.scaling = scaling
+ self.head_dropout = head_dropout
+ self.num_patches = (max(context_length, patch_length) - patch_length) // patch_stride + 1
+ self.mask_type = mask_type
+ self.random_mask_ratio = random_mask_ratio
+ self.num_forecast_mask_patches = num_forecast_mask_patches
+ self.mask_value = mask_value
+ self.channel_consistent_masking = channel_consistent_masking
+ self.masked_loss = masked_loss
+ self.patch_last = True
+ self.use_positional_encoding = use_positional_encoding
+ self.positional_encoding_type = positional_encoding_type
+ self.prediction_length = prediction_length
+ self.prediction_channel_indices = prediction_channel_indices
+ self.num_targets = num_targets
+ self.output_range = output_range
+ self.head_aggregation = head_aggregation
+ self.self_attn = self_attn
+ self.self_attn_heads = self_attn_heads
+ self.init_std = init_std
+ self.post_init = post_init
+ self.distribution_output = distribution_output
+ self.loss = loss
+ self.num_parallel_samples = num_parallel_samples
+ self.unmasked_channel_indices = unmasked_channel_indices
+ self.norm_eps = norm_eps
+ super().__init__(**kwargs)
+
+
+__all__ = ["PatchTSMixerConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/patchtsmixer/modeling_patchtsmixer.py b/lib/python3.12/site-packages/transformers/models/patchtsmixer/modeling_patchtsmixer.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca88a84a399ef11729898ed22c5fc884e9c23e70
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/patchtsmixer/modeling_patchtsmixer.py
@@ -0,0 +1,2182 @@
+# coding=utf-8
+# Copyright 2023 IBM 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 PatchTSMixer model."""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from .configuration_patchtsmixer import PatchTSMixerConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "PatchTSMixerConfig"
+
+
+PATCHTSMIXER_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 ([`PatchTSMixerConfig`]):
+ 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.
+ mask_input (`bool`, *optional*, defaults to `False`):
+ If True, Masking will be enabled. False otherwise.
+"""
+
+PATCHTSMIXER_INPUTS_DOCSTRING = r"""
+ Args:
+ past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`):
+ Context values of the time series. For a pretraining task, this denotes the input time series to predict
+ the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly,
+ for classification or regression tasks, it denotes the appropriate context values of the time series.
+
+ For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is
+ greater than 1.
+
+ 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.
+"""
+
+
+class PatchTSMixerGatedAttention(nn.Module):
+ """
+ Module that applies gated attention to input data.
+
+ Args:
+ in_size (`int`): The input size.
+ out_size (`int`): The output size.
+ """
+
+ def __init__(self, in_size: int, out_size: int):
+ super().__init__()
+ self.attn_layer = nn.Linear(in_size, out_size)
+ self.attn_softmax = nn.Softmax(dim=-1)
+
+ def forward(self, inputs):
+ attn_weight = self.attn_softmax(self.attn_layer(inputs))
+ inputs = inputs * attn_weight
+ return inputs
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTBatchNorm with PatchTST->PatchTSMixer
+class PatchTSMixerBatchNorm(nn.Module):
+ """
+ Compute batch normalization over the sequence length (time) dimension.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ self.batchnorm = nn.BatchNorm1d(config.d_model, eps=config.norm_eps)
+
+ def forward(self, inputs: torch.Tensor):
+ """
+ Parameters:
+ inputs (`torch.Tensor` of shape `(batch_size, sequence_length, d_model)`):
+ input for Batch norm calculation
+ Returns:
+ `torch.Tensor` of shape `(batch_size, sequence_length, d_model)`
+ """
+ output = inputs.transpose(1, 2) # output: (batch_size, d_model, sequence_length)
+ output = self.batchnorm(output)
+ return output.transpose(1, 2)
+
+
+class PatchTSMixerPositionalEncoding(nn.Module):
+ """
+ Class for positional encoding
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ # positional encoding: [num_patches x d_model]
+ if config.use_positional_encoding:
+ self.position_enc = self._init_pe(config)
+ else:
+ self.position_enc = nn.Parameter(torch.zeros(config.num_patches, config.d_model))
+
+ @staticmethod
+ def _init_pe(config: PatchTSMixerConfig) -> nn.Parameter:
+ # Positional encoding
+ if config.positional_encoding_type == "random":
+ position_enc = nn.Parameter(torch.randn(config.num_patches, config.d_model), requires_grad=True)
+ elif config.positional_encoding_type == "sincos":
+ position_enc = torch.zeros(config.num_patches, config.d_model)
+ position = torch.arange(0, config.num_patches).unsqueeze(1)
+ div_term = torch.exp(torch.arange(0, config.d_model, 2) * -(math.log(10000.0) / config.d_model))
+ position_enc[:, 0::2] = torch.sin(position * div_term)
+ position_enc[:, 1::2] = torch.cos(position * div_term)
+ position_enc = position_enc - position_enc.mean()
+ position_enc = position_enc / (position_enc.std() * 10)
+ position_enc = nn.Parameter(position_enc, requires_grad=False)
+ else:
+ raise ValueError(
+ f"{config.positional_encoding_type} is not a valid positional encoder. Available types are 'random' and 'sincos'."
+ )
+ return position_enc
+
+ def forward(self, patch_input: torch.Tensor):
+ # hidden_state: [bs x num_channels x num_patches x d_model]
+ hidden_state = patch_input + self.position_enc
+ return hidden_state
+
+
+class PatchTSMixerNormLayer(nn.Module):
+ """Normalization block
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.norm_mlp = config.norm_mlp
+
+ if "batch" in config.norm_mlp.lower():
+ self.norm = PatchTSMixerBatchNorm(config)
+ else:
+ self.norm = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+
+ def forward(self, inputs: torch.Tensor):
+ """
+ Args:
+ inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`):
+ Input to the normalization layer.
+ Returns:
+ `torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`
+ """
+ if "batch" in self.norm_mlp.lower():
+ # reshape the data
+ inputs_reshaped = torch.reshape(
+ inputs,
+ (
+ inputs.shape[0] * inputs.shape[1],
+ inputs.shape[2],
+ inputs.shape[3],
+ ),
+ ) # inputs_reshaped: [batch_size*num_channels, num_patches, d_model]
+
+ # inputs_reshaped: [batch_size*num_channels, num_patches, d_model]
+ inputs_reshaped = self.norm(inputs_reshaped)
+
+ # put back data to the original shape
+ inputs = torch.reshape(inputs_reshaped, inputs.shape)
+
+ else:
+ inputs = self.norm(inputs)
+
+ return inputs
+
+
+class PatchTSMixerMLP(nn.Module):
+ def __init__(self, in_features, out_features, config):
+ super().__init__()
+ num_hidden = in_features * config.expansion_factor
+ self.fc1 = nn.Linear(in_features, num_hidden)
+ self.dropout1 = nn.Dropout(config.dropout)
+ self.fc2 = nn.Linear(num_hidden, out_features)
+ self.dropout2 = nn.Dropout(config.dropout)
+
+ def forward(self, inputs: torch.Tensor):
+ """
+ Args:
+ inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`):
+ Input to the MLP layer.
+ Returns:
+ `torch.Tensor` of the same shape as `inputs`
+ """
+ inputs = self.dropout1(nn.functional.gelu(self.fc1(inputs)))
+ inputs = self.fc2(inputs)
+ inputs = self.dropout2(inputs)
+ return inputs
+
+
+class PatchTSMixerChannelFeatureMixerBlock(nn.Module):
+ """This module mixes the features in the channel dimension.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.norm = PatchTSMixerNormLayer(config)
+ self.gated_attn = config.gated_attn
+ self.mlp = PatchTSMixerMLP(
+ in_features=config.num_input_channels,
+ out_features=config.num_input_channels,
+ config=config,
+ )
+
+ if config.gated_attn:
+ self.gating_block = PatchTSMixerGatedAttention(
+ in_size=config.num_input_channels, out_size=config.num_input_channels
+ )
+
+ def forward(self, inputs: torch.Tensor):
+ """
+ Args:
+ inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`):
+ input to the MLP layer
+ Returns:
+ `torch.Tensor` of the same shape as `inputs`
+ """
+ residual = inputs
+ inputs = self.norm(inputs)
+
+ inputs = inputs.permute(0, 3, 2, 1)
+
+ if self.gated_attn:
+ inputs = self.gating_block(inputs)
+
+ inputs = self.mlp(inputs)
+
+ inputs = inputs.permute(0, 3, 2, 1)
+
+ out = inputs + residual
+ return out
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->PatchTSMixer
+class PatchTSMixerAttention(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[PatchTSMixerConfig] = 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 PatchMixerBlock(nn.Module):
+ """This module mixes the patch dimension.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.norm = PatchTSMixerNormLayer(config)
+
+ self.self_attn = config.self_attn
+ self.gated_attn = config.gated_attn
+
+ self.mlp = PatchTSMixerMLP(
+ in_features=config.num_patches,
+ out_features=config.num_patches,
+ config=config,
+ )
+
+ if config.gated_attn:
+ self.gating_block = PatchTSMixerGatedAttention(in_size=config.num_patches, out_size=config.num_patches)
+
+ if config.self_attn:
+ self.self_attn_layer = PatchTSMixerAttention(
+ embed_dim=config.d_model,
+ num_heads=config.self_attn_heads,
+ dropout=config.dropout,
+ )
+ self.norm_attn = PatchTSMixerNormLayer(config)
+
+ def forward(self, hidden_state):
+ """
+ Args:
+ hidden_state (`torch.Tensor`): Input tensor.
+
+ Returns:
+ `torch.Tensor`: Transformed tensor.
+ """
+ residual = hidden_state
+
+ hidden_state = self.norm(hidden_state)
+
+ if self.self_attn:
+ batch_size, n_vars, num_patches, d_model = hidden_state.shape
+ hidden_state_reshaped = hidden_state.reshape(batch_size * n_vars, num_patches, d_model)
+
+ x_attn, _, _ = self.self_attn_layer(hidden_state_reshaped, output_attentions=False)
+ x_attn = x_attn.reshape(batch_size, n_vars, num_patches, d_model)
+
+ # Transpose so that num_patches is the last dimension
+ hidden_state = hidden_state.transpose(2, 3)
+ hidden_state = self.mlp(hidden_state)
+
+ if self.gated_attn:
+ hidden_state = self.gating_block(hidden_state)
+
+ # Transpose back
+ hidden_state = hidden_state.transpose(2, 3)
+
+ if self.self_attn:
+ hidden_state = self.norm_attn(hidden_state + x_attn)
+
+ out = hidden_state + residual
+ return out
+
+
+class FeatureMixerBlock(nn.Module):
+ """This module mixes the hidden feature dimension.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.norm = PatchTSMixerNormLayer(config)
+
+ self.gated_attn = config.gated_attn
+
+ self.mlp = PatchTSMixerMLP(
+ in_features=config.d_model,
+ out_features=config.d_model,
+ config=config,
+ )
+
+ if config.gated_attn:
+ self.gating_block = PatchTSMixerGatedAttention(in_size=config.d_model, out_size=config.d_model)
+
+ def forward(self, hidden: torch.Tensor):
+ """
+ Args:
+ hidden (`torch.Tensor` of shape `(batch_size, num_patches, d_model)`):
+ Input tensor to the layer.
+
+ Returns:
+ `torch.Tensor`: Transformed tensor.
+ """
+ residual = hidden
+ hidden = self.norm(hidden)
+ hidden = self.mlp(hidden)
+
+ if self.gated_attn:
+ hidden = self.gating_block(hidden)
+
+ out = hidden + residual
+ return out
+
+
+class PatchTSMixerLayer(nn.Module):
+ """
+ The `PatchTSMixer` layer that does all three kinds of mixing.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.patch_mixer = PatchMixerBlock(config=config)
+ self.feature_mixer = FeatureMixerBlock(config=config)
+
+ self.mode = config.mode
+
+ if config.mode == "mix_channel":
+ self.channel_feature_mixer = PatchTSMixerChannelFeatureMixerBlock(config=config)
+
+ def forward(self, hidden: torch.Tensor):
+ """
+ Args:
+ hidden (`torch.Tensor` of shape `(batch_size, num_patches, d_model)`):
+ Input tensor to the layer.
+
+ Returns:
+ `torch.Tensor`: Transformed tensor.
+ """
+ if self.mode == "mix_channel":
+ hidden = self.channel_feature_mixer(hidden)
+
+ hidden = self.patch_mixer(hidden)
+ hidden = self.feature_mixer(hidden) # hidden: (batch_size x num_patches x d_model)
+ return hidden
+
+
+class PatchTSMixerBlock(nn.Module):
+ """The main computing framework of the `PatchTSMixer` model.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ num_layers = config.num_layers
+
+ self.mixers = nn.ModuleList([PatchTSMixerLayer(config=config) for _ in range(num_layers)])
+
+ def forward(self, hidden_state, output_hidden_states: bool = False):
+ """
+ Args:
+ hidden_state (`torch.Tensor`): The input tensor.
+ output_hidden_states (`bool`, *optional*, defaults to False.):
+ Whether to output the hidden states as well.
+
+ Returns:
+ `torch.Tensor`: The embedding. `list`: List of all hidden states if `output_hidden_states` is set to
+ `True`.
+ """
+ all_hidden_states = []
+
+ embedding = hidden_state
+
+ for mod in self.mixers:
+ embedding = mod(embedding)
+ if output_hidden_states:
+ all_hidden_states.append(embedding)
+
+ if output_hidden_states:
+ return embedding, all_hidden_states
+ else:
+ return embedding, None
+
+
+class PatchTSMixerForPredictionHead(nn.Module):
+ """Prediction Head for Forecasting
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig, distribution_output=None):
+ super().__init__()
+
+ self.prediction_channel_indices = config.prediction_channel_indices
+
+ if self.prediction_channel_indices is not None:
+ self.prediction_channel_indices.sort()
+
+ self.dropout_layer = nn.Dropout(config.head_dropout)
+ if distribution_output is None:
+ self.base_forecast_block = nn.Linear((config.num_patches * config.d_model), config.prediction_length)
+ else:
+ self.base_forecast_block = distribution_output.get_parameter_projection(
+ config.num_patches * config.d_model
+ )
+
+ self.flatten = nn.Flatten(start_dim=-2)
+
+ def forward(self, hidden_features):
+ """
+
+ Args:
+ hidden_features (`torch.Tensor` of shape `(batch_size, num_patch, d_model)` in `flatten` mode
+ or `(batch_size, n_vars, num_patch, d_model)` in `common_channel`/`mix_channel` mode.): Input hidden
+ features.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size, prediction_length, nvars)`.
+
+ """
+
+ hidden_features = self.flatten(hidden_features) # [batch_size x n_vars x num_patch * d_model]
+ hidden_features = self.dropout_layer(hidden_features) # [batch_size x n_vars x num_patch * d_model]
+ forecast = self.base_forecast_block(hidden_features) # [batch_size x n_vars x prediction_length]
+ if isinstance(forecast, tuple):
+ forecast = tuple(z.transpose(-1, -2) for z in forecast)
+ else:
+ forecast = forecast.transpose(-1, -2) # [batch_size x prediction_length x n_vars]
+
+ if self.prediction_channel_indices is not None:
+ if isinstance(forecast, tuple):
+ forecast = tuple(z[..., self.prediction_channel_indices] for z in forecast)
+ else:
+ forecast = forecast[..., self.prediction_channel_indices] # [batch_size x prediction_length x n_vars]
+
+ return forecast
+
+
+class PatchTSMixerLinearHead(nn.Module):
+ """Linear head for Classification and Regression.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig, distribution_output=None):
+ super().__init__()
+
+ self.head_aggregation = config.head_aggregation
+ self.output_range = config.output_range
+
+ if config.head_aggregation is None:
+ mul_factor = config.num_patches
+ else:
+ mul_factor = 1
+ self.distribution_output = distribution_output
+ if distribution_output is None:
+ self.projection = nn.Linear(
+ config.d_model * config.num_input_channels * mul_factor,
+ config.num_targets,
+ )
+ else:
+ self.projection = distribution_output.get_parameter_projection(
+ config.d_model * config.num_input_channels * mul_factor
+ )
+
+ if config.head_aggregation is None:
+ self.flatten = nn.Flatten(start_dim=-3)
+ else:
+ self.flatten = nn.Flatten(start_dim=-2)
+
+ self.dropout = nn.Dropout(config.head_dropout)
+
+ def forward(self, hidden_features):
+ """
+ Args:
+ hidden_features (`torch.Tensor` of shape `(batch_size x num_patch x d_model)` in `flatten` mode
+ or `(batch_size x n_vars x num_patch x d_model)` in `common_channel`/`mix_channel` mode.): Input hidden
+ features.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size x num_targets)`.
+ """
+
+ # batch_size x d_model x num_patch or batch_size x n_vars x d_model x num_patch
+ hidden_features = hidden_features.transpose(-1, -2)
+ if self.head_aggregation == "use_last":
+ # batch_size x d_model (flatten) or # batch_size x n_vars x d_model (common_channel)
+ hidden_features = hidden_features[..., -1]
+ elif self.head_aggregation == "max_pool":
+ # batch_size x n_vars x d_model or batch_size x d_model
+ hidden_features = hidden_features.max(dim=-1).values
+ elif self.head_aggregation == "avg_pool":
+ # batch_size x n_vars x d_model or batch_size x d_model
+ hidden_features = hidden_features.mean(dim=-1)
+
+ if self.flatten:
+ hidden_features = self.flatten(hidden_features)
+ hidden_features = self.dropout(hidden_features)
+ hidden_features = self.projection(hidden_features) # batch_size x num_targets
+
+ if (self.distribution_output is None) and (self.output_range is not None):
+ hidden_features = (
+ torch.sigmoid(hidden_features) * (self.output_range[1] - self.output_range[0]) + self.output_range[0]
+ )
+ return hidden_features
+
+
+class PatchTSMixerPreTrainedModel(PreTrainedModel):
+ # Weight initialization
+ config_class = PatchTSMixerConfig
+ base_model_prefix = "model"
+ main_input_name = "past_values"
+ supports_gradient_checkpointing = False
+
+ def _init_weights(self, module):
+ """Initialize weights"""
+ if isinstance(module, PatchTSMixerPositionalEncoding):
+ # initialize positional encoding
+ if self.config.positional_encoding_type == "random":
+ nn.init.normal_(module.position_enc, mean=0.0, std=0.1)
+ elif isinstance(module, (nn.LayerNorm, nn.BatchNorm1d)):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+ elif isinstance(module, PatchTSMixerBatchNorm):
+ module.batchnorm.bias.data.zero_()
+ module.batchnorm.weight.data.fill_(1.0)
+ elif 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_()
+
+
+class PatchTSMixerPretrainHead(nn.Module):
+ """Pretraining head.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.dropout_layer = nn.Dropout(config.head_dropout)
+ self.base_pt_block = nn.Linear(config.d_model, config.patch_length)
+
+ def forward(self, hidden_features):
+ """
+ Args:
+ hidden_features (`torch.Tensor` of shape `(batch_size x num_patch x d_model)` in `flatten` mode
+ or `(batch_size x n_vars x num_patch x d_model)` in `common_channel`/`mix_channel` mode.): Input hidden
+ features.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size x n_vars x num_patch x patch_length)`.
+ """
+
+ hidden_features = self.dropout_layer(hidden_features)
+ forecast = self.base_pt_block(hidden_features) # [batch_size x n_vars x num_patch x patch_length]
+ return forecast
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.random_masking
+def random_masking(
+ inputs: torch.Tensor,
+ mask_ratio: float,
+ unmasked_channel_indices: list = None,
+ channel_consistent_masking: bool = False,
+ mask_value: int = 0,
+):
+ """random_masking: Mask the input considering the control variables.
+
+ Args:
+ inputs (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, num_features)`):
+ The input tensor to mask.
+ mask_ratio (`float`):
+ Masking ratio applied to mask the input data during random pretraining. It is the number between 0 and 1.
+ unmasked_channel_indices (list, *optional*):
+ Indices of channels that will not be masked.
+ channel_consistent_masking (bool, *optional*, defaults to `False`):
+ When true, masking will be same across all channels of a timeseries. Otherwise, masking positions will vary
+ across channels.
+ mask_value (int, *optional*, defaults to 0):
+ Define the value of masked patches for pretraining.
+
+ Returns:
+ `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as input Tensor and mask tensor of shape [bs x c x
+ n]
+ """
+ if mask_ratio < 0 or mask_ratio >= 1:
+ raise ValueError(f"Mask ratio {mask_ratio} has to be between 0 and 1.")
+
+ batch_size, num_channels, sequence_length, num_features = inputs.shape
+ device = inputs.device
+
+ len_keep = int(sequence_length * (1 - mask_ratio))
+
+ if channel_consistent_masking:
+ noise = torch.rand(batch_size, 1, sequence_length, device=device) # noise in [0, 1], bs x 1 x L
+ noise = noise.repeat(1, num_channels, 1) # bs x num_channels x time
+ else:
+ # noise in [0, 1], bs x num_channels x L
+ noise = torch.rand(batch_size, num_channels, sequence_length, device=device)
+
+ # mask: [bs x num_channels x num_patch]
+ mask = torch.ones(batch_size, num_channels, sequence_length, device=device)
+ mask[:, :, :len_keep] = 0
+
+ # sort noise for each sample
+ ids_shuffle = torch.argsort(noise, dim=-1) # ascend: small is keep, large is remove
+ ids_restore = torch.argsort(ids_shuffle, dim=-1) # ids_restore: [bs x num_channels x L]
+
+ mask = torch.gather(mask, dim=-1, index=ids_restore)
+ mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features) # mask: [bs x num_channels x num_patches x patch_length]
+ if unmasked_channel_indices is not None:
+ mask[:, unmasked_channel_indices, :, :] = 0
+
+ inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
+ return inputs_mask, mask[..., 0]
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.forecast_masking
+def forecast_masking(
+ inputs: torch.Tensor,
+ num_forecast_mask_patches: Union[list, int],
+ unmasked_channel_indices: list = None,
+ mask_value: int = 0,
+):
+ """Forecast masking that masks the last K patches where K is from the num_forecast_mask_patches.
+ If num_forecast_mask_patches is a list, samples in the batch will be randomly masked by numbers defined in the list.
+
+ Parameters:
+ inputs (`torch.Tensor`):
+ Input of shape `(bs, num_channels, num_patch, patch_length)`
+ num_forecast_mask_patches (`list`):
+ Number of patches to be masked at the end of each batch sample. e.g. 4 or [3, 5].
+ unmasked_channel_indices (`list`, *optional*):
+ Indices of channels that are not masked.
+ mask_value (`int`, *optional*, defaults to 0):
+ Values in the masked patches will be filled by `mask_value`.
+
+ Returns:
+ `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as inputs Tensor and Mask tensor of shape `(bs,
+ num_channels , num_patch)` or `(bs, tsg1, tsg2, num_channels, num_patch)`
+ """
+
+ if isinstance(num_forecast_mask_patches, int):
+ num_forecast_mask_patches = [num_forecast_mask_patches]
+ forecast_mask_ratios = [1 for _ in num_forecast_mask_patches]
+
+ batch_size, num_channels, sequence_length, num_features = inputs.shape
+ mask = torch.zeros(batch_size, num_channels, sequence_length, device=inputs.device)
+
+ t_list = []
+ total_length = 0
+ total_ratio = sum(forecast_mask_ratios)
+
+ for patch_length, ratio in zip(num_forecast_mask_patches, forecast_mask_ratios):
+ if patch_length <= 0 or patch_length >= sequence_length:
+ raise ValueError(
+ f"num_forecast_mask_patches {patch_length} should be greater than 0 and less than total patches."
+ )
+ temp_len = int(batch_size * ratio / total_ratio)
+ t_list.append([patch_length, ratio, temp_len])
+ total_length += temp_len
+
+ t_list = sorted(t_list, key=lambda x: x[2])
+
+ if total_length < batch_size:
+ t_list[0][2] = t_list[0][2] + (batch_size - total_length)
+ elif total_length > batch_size:
+ t_list[-1][2] = t_list[-1][2] + (total_length - batch_size)
+
+ batch1 = 0
+ for patch_len, _, temp_len in t_list:
+ batch2 = batch1 + temp_len
+ mask[batch1:batch2, :, -patch_len:] = 1
+ batch1 = batch2
+
+ perm = torch.randperm(mask.shape[0])
+ mask = mask[perm]
+
+ mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features) # mask: [bs x num_channels x num_patch x patch_len]
+ if unmasked_channel_indices is not None:
+ mask[:, unmasked_channel_indices, :, :] = 0
+
+ inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
+ return inputs_mask, mask[..., 0]
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTPatchify with PatchTST->PatchTSMixer
+class PatchTSMixerPatchify(nn.Module):
+ """
+ A class to patchify the time series sequence into different patches
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+
+ self.sequence_length = config.context_length
+ self.patch_length = config.patch_length
+ self.patch_stride = config.patch_stride
+
+ if self.sequence_length <= self.patch_length:
+ raise ValueError(
+ f"Sequence length ({self.sequence_length}) has to be greater than the patch length ({self.patch_length})"
+ )
+
+ # get the number of patches
+ self.num_patches = (max(self.sequence_length, self.patch_length) - self.patch_length) // self.patch_stride + 1
+ new_sequence_length = self.patch_length + self.patch_stride * (self.num_patches - 1)
+ self.sequence_start = self.sequence_length - new_sequence_length
+
+ def forward(self, past_values: torch.Tensor):
+ """
+ Parameters:
+ past_values (`torch.Tensor` of shape `(batch_size, sequence_length, num_channels)`, *required*):
+ Input for patchification
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`
+ """
+ sequence_length = past_values.shape[-2]
+ if sequence_length != self.sequence_length:
+ raise ValueError(
+ f"Input sequence length ({sequence_length}) doesn't match model configuration ({self.sequence_length})."
+ )
+ # output: [bs x new_sequence_length x num_channels]
+ output = past_values[:, self.sequence_start :, :]
+ # output: [bs x num_patches x num_input_channels x patch_length]
+ output = output.unfold(dimension=-2, size=self.patch_length, step=self.patch_stride)
+ # output: [bs x num_input_channels x num_patches x patch_length]
+ output = output.transpose(-2, -3).contiguous()
+ return output
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTMasking with PatchTST->PatchTSMixer
+class PatchTSMixerMasking(nn.Module):
+ """
+ Class to perform random or forecast masking.
+
+ Parameters:
+ config (`PatchTSMixerConfig`): model config
+ Returns:
+ x_mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)
+ Masked patched input
+ mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)
+ Bool tensor indicating True on masked points
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ self.random_mask_ratio = config.random_mask_ratio
+ self.channel_consistent_masking = config.channel_consistent_masking
+ self.mask_type = config.mask_type
+ self.num_forecast_mask_patches = config.num_forecast_mask_patches
+ self.unmasked_channel_indices = config.unmasked_channel_indices
+ self.mask_value = config.mask_value
+ if self.unmasked_channel_indices is not None:
+ self.unmasked_channel_indices = sorted(self.unmasked_channel_indices)
+
+ def forward(self, patch_input: torch.Tensor):
+ """
+ Parameters:
+ patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):
+ Patch input
+
+ Return:
+ masked_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)
+ Masked patched input
+ mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)
+ Bool tensor indicating True on masked points
+
+ """
+ if self.mask_type == "random":
+ masked_input, mask = random_masking(
+ inputs=patch_input,
+ mask_ratio=self.random_mask_ratio,
+ unmasked_channel_indices=self.unmasked_channel_indices,
+ channel_consistent_masking=self.channel_consistent_masking,
+ mask_value=self.mask_value,
+ )
+ elif self.mask_type == "forecast":
+ masked_input, mask = forecast_masking(
+ inputs=patch_input,
+ num_forecast_mask_patches=self.num_forecast_mask_patches,
+ unmasked_channel_indices=self.unmasked_channel_indices,
+ mask_value=self.mask_value,
+ )
+ else:
+ raise ValueError(f"Invalid mask type {self.mask_type}.")
+
+ # mask: [bs x num_input_channels x num_patch]
+ mask = mask.bool()
+ return masked_input, mask
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTStdScaler with PatchTST->PatchTSMixer
+class PatchTSMixerStdScaler(nn.Module):
+ """
+ Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by
+ subtracting from the mean and dividing by the standard deviation.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+ self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Calculating the scale on the observed indicator.
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)
+ denominator = denominator.clamp_min(1.0)
+ loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator
+
+ variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator
+ scale = torch.sqrt(variance + self.minimum_scale)
+ return (data - loc) / scale, loc, scale
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTMeanScaler with PatchTST->PatchTSMixer
+class PatchTSMixerMeanScaler(nn.Module):
+ """
+ Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data
+ accordingly.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+ self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10
+ self.default_scale = config.default_scale if hasattr(config, "default_scale") else None
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Calculating the scale on the observed indicator.
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
+ num_observed = observed_indicator.sum(self.dim, keepdim=True)
+
+ scale = ts_sum / torch.clamp(num_observed, min=1)
+
+ # If `default_scale` is provided, we use it, otherwise we use the scale
+ # of the batch.
+ if self.default_scale is None:
+ batch_sum = ts_sum.sum(dim=0)
+ batch_observations = torch.clamp(num_observed.sum(0), min=1)
+ default_scale = torch.squeeze(batch_sum / batch_observations)
+ else:
+ default_scale = self.default_scale * torch.ones_like(scale)
+
+ # apply default scale where there are no observations
+ scale = torch.where(num_observed > 0, scale, default_scale)
+
+ # ensure the scale is at least `self.minimum_scale`
+ scale = torch.clamp(scale, min=self.minimum_scale)
+ scaled_data = data / scale
+
+ if not self.keepdim:
+ scale = scale.squeeze(dim=self.dim)
+
+ return scaled_data, torch.zeros_like(scale), scale
+
+
+# Copied from transformers.models.patchtst.modeling_patchtst.PatchTSTNOPScaler with PatchTST->PatchTSMixer
+class PatchTSMixerNOPScaler(nn.Module):
+ """
+ Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: Optional[torch.Tensor] = None
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+ loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+ return data, loc, scale
+
+
+@dataclass
+class PatchTSMixerEncoderOutput(ModelOutput):
+ """
+ Base class for `PatchTSMixerEncoderOutput`, with potential hidden states.
+
+ Args:
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, d_model)`):
+ Hidden-state at the output of the last layer of the model.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class PatchTSMixerEncoder(PatchTSMixerPreTrainedModel):
+ """
+ Encoder for PatchTSMixer which inputs patched time-series and outputs patched embeddings.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__(config)
+
+ self.use_return_dict = config.use_return_dict
+
+ self.patcher = nn.Linear(config.patch_length, config.d_model)
+ if config.use_positional_encoding:
+ self.positional_encoder = PatchTSMixerPositionalEncoding(config=config)
+ else:
+ self.positional_encoder = None
+ self.mlp_mixer_encoder = PatchTSMixerBlock(config=config)
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @replace_return_docstrings(output_type=PatchTSMixerEncoderOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ output_hidden_states: Optional[bool] = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, PatchTSMixerEncoderOutput]:
+ r"""
+ Args:
+ past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`):
+ Context values of the time series. For a pretraining task, this denotes the input time series to
+ predict the masked portion. For a forecasting task, this denotes the history/past time series values.
+ Similarly, for classification or regression tasks, it denotes the appropriate context values of the
+ time series.
+
+ For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series,
+ it is greater than 1.
+
+ 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.
+
+ Returns:
+ `torch.FloatTensor` of shape `(batch_size, n_vars, num_patches, d_model)`
+ """
+
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+
+ # flatten [bs x num_patch x d_model]. common_channel/mix_channel: [bs x n_vars x num_patch x d_model]
+ patches = self.patcher(past_values)
+
+ # add positional encoder
+ if self.positional_encoder is not None:
+ patches = self.positional_encoder(patches)
+
+ last_hidden_state, hidden_states = self.mlp_mixer_encoder(patches, output_hidden_states=output_hidden_states)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ last_hidden_state,
+ hidden_states,
+ ]
+ )
+
+ return PatchTSMixerEncoderOutput(last_hidden_state=last_hidden_state, hidden_states=hidden_states)
+
+
+@dataclass
+class PatchTSMixerModelOutput(ModelOutput):
+ """
+ Base class for model's outputs, with potential hidden states.
+
+ Args:
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, d_model)`):
+ Hidden-state at the output of the last layer of the model.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer.
+ patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`):
+ Patched input data to the model.
+ mask: (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`,*optional*):
+ Bool Tensor indicating True in masked patches and False otherwise.
+ loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`,*optional*):
+ Gives the mean of the context window per channel. Used for revin denorm outside the model, if revin
+ enabled.
+ scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`,*optional*):
+ Gives the std dev of the context window per channel. Used for revin denorm outside the model, if revin
+ enabled.
+ """
+
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ patch_input: Optional[torch.FloatTensor] = None
+ mask: Optional[torch.FloatTensor] = None
+ loc: Optional[torch.FloatTensor] = None
+ scale: Optional[torch.FloatTensor] = None
+
+
+@add_start_docstrings(
+ "The PatchTSMixer Model for time-series forecasting.",
+ PATCHTSMIXER_START_DOCSTRING,
+)
+class PatchTSMixerModel(PatchTSMixerPreTrainedModel):
+ def __init__(self, config: PatchTSMixerConfig, mask_input: bool = False):
+ super().__init__(config)
+
+ self.use_return_dict = config.use_return_dict
+ self.encoder = PatchTSMixerEncoder(config)
+ self.patching = PatchTSMixerPatchify(config)
+
+ if mask_input is True:
+ self.masking = PatchTSMixerMasking(config)
+ else:
+ self.masking = None
+
+ if config.scaling == "mean":
+ self.scaler = PatchTSMixerMeanScaler(config)
+ elif config.scaling == "std" or config.scaling is True:
+ self.scaler = PatchTSMixerStdScaler(config)
+ else:
+ self.scaler = PatchTSMixerNOPScaler(config)
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(PATCHTSMIXER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=PatchTSMixerModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ observed_mask: Optional[torch.Tensor] = None,
+ output_hidden_states: Optional[bool] = False,
+ return_dict: Optional[bool] = None,
+ ) -> PatchTSMixerModelOutput:
+ r"""
+ observed_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+ Returns:
+
+ """
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+
+ mask = None
+ if observed_mask is None:
+ observed_mask = torch.ones_like(past_values)
+ scaled_past_values, loc, scale = self.scaler(past_values, observed_mask)
+
+ patched_x = self.patching(scaled_past_values) # [batch_size x num_input_channels x num_patch x patch_length
+
+ enc_input = patched_x
+ if self.masking is not None:
+ enc_input, mask = self.masking(patched_x)
+ # enc_input: [batch_size x num_input_channels x num_patch x patch_length]
+ # mask: [batch_size x num_input_channels x num_patch]
+
+ encoder_output = self.encoder(
+ enc_input,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ if isinstance(encoder_output, tuple):
+ encoder_output = PatchTSMixerEncoderOutput(*encoder_output)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ encoder_output.last_hidden_state,
+ encoder_output.hidden_states,
+ patched_x,
+ mask,
+ loc,
+ scale,
+ ]
+ )
+
+ return PatchTSMixerModelOutput(
+ last_hidden_state=encoder_output.last_hidden_state,
+ hidden_states=encoder_output.hidden_states,
+ patch_input=patched_x,
+ mask=mask,
+ loc=loc,
+ scale=scale,
+ )
+
+
+@dataclass
+class PatchTSMixerForPreTrainingOutput(ModelOutput):
+ """
+ Output type of [`PatchTSMixerForPreTrainingOutput`].
+
+ Args:
+ prediction_outputs (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, patch_length)`):
+ Prediction output from the pretrain head.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`):
+ Backbone embeddings before passing through the head.
+ loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`):
+ Total loss
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ prediction_outputs: Optional[torch.FloatTensor] = None
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class PatchTSMixerForPretraining(PatchTSMixerPreTrainedModel):
+ r"""
+ `PatchTSMixer` for mask pretraining.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ Returns:
+ `None`.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__(config)
+ self.model = PatchTSMixerModel(config, mask_input=True)
+ self.head = PatchTSMixerPretrainHead(config=config)
+ self.masked_loss = config.masked_loss
+ self.use_return_dict = config.use_return_dict
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(PATCHTSMIXER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=PatchTSMixerForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ observed_mask: Optional[torch.Tensor] = None,
+ output_hidden_states: Optional[bool] = False,
+ return_loss: bool = True,
+ return_dict: Optional[bool] = None,
+ ) -> PatchTSMixerForPreTrainingOutput:
+ r"""
+ observed_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+ return_loss (`bool`, *optional*):
+ Whether to return the loss in the `forward` call.
+
+ Returns:
+
+ """
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+
+ if self.masked_loss is True:
+ loss = torch.nn.MSELoss(reduction="none")
+ else:
+ loss = torch.nn.MSELoss(reduction="mean")
+
+ # past_values: tensor [batch_size x context_length x num_input_channels]
+ model_output = self.model(
+ past_values,
+ observed_mask=observed_mask,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ ) # x.last_hidden_state: [batch_size x nvars x num_patch x d_model]
+ if isinstance(model_output, tuple):
+ model_output = PatchTSMixerModelOutput(*model_output)
+
+ x_hat = self.head(model_output.last_hidden_state) # tensor [batch_size x nvars x num_patch x patch_length]
+
+ if return_loss is True:
+ loss_val = loss(x_hat, model_output.patch_input)
+ else:
+ loss_val = None
+
+ # calculate masked_loss
+ if self.masked_loss is True and loss_val is not None:
+ loss_val = (loss_val.mean(dim=-1) * model_output.mask).sum() / (model_output.mask.sum() + 1e-10)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ loss_val,
+ x_hat,
+ model_output.last_hidden_state,
+ model_output.hidden_states,
+ ]
+ )
+
+ return PatchTSMixerForPreTrainingOutput(
+ loss=loss_val,
+ prediction_outputs=x_hat, # tensor [batch_size x nvars x num_patch x patch_length]
+ last_hidden_state=model_output.last_hidden_state, # x: [batch_size x nvars x num_patch x d_model]
+ hidden_states=model_output.hidden_states,
+ )
+
+
+@dataclass
+class PatchTSMixerForPredictionOutput(ModelOutput):
+ """
+ Output type of [`PatchTSMixerForPredictionOutput`].
+
+ Args:
+ prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_input_channels)`):
+ Prediction output from the forecast head.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`):
+ Backbone embeddings before passing through the head.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+ loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`):
+ Total loss.
+ loc (`torch.FloatTensor`, *optional* of shape `(batch_size, 1, num_input_channels)`):
+ Input mean
+ scale (`torch.FloatTensor`, *optional* of shape `(batch_size, 1, num_input_channels)`):
+ Input std dev
+
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ prediction_outputs: Optional[torch.FloatTensor] = None
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ loc: Optional[torch.FloatTensor] = None
+ scale: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class SamplePatchTSMixerPredictionOutput(ModelOutput):
+ """
+ Base class for time series model's predictions outputs that contains the sampled values from the chosen
+ distribution.
+
+ Args:
+ sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, number_channels)`):
+ Sampled values from the chosen distribution.
+ """
+
+ sequences: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class SamplePatchTSMixerRegressionOutput(ModelOutput):
+ """
+ Base class for time series model's predictions outputs that contains the sampled values from the chosen
+ distribution.
+
+ Args:
+ sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, num_targets)`
+ Sampled values from the chosen distribution.
+ """
+
+ sequences: Optional[torch.FloatTensor] = None
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.nll
+def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor:
+ """
+ Computes the negative log likelihood loss from input distribution with respect to target.
+ """
+ return -input.log_prob(target)
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.weighted_average
+def weighted_average(input_tensor: torch.Tensor, weights: Optional[torch.Tensor] = None, dim=None) -> torch.Tensor:
+ """
+ Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero,
+ meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.
+
+ Args:
+ input_tensor (`torch.FloatTensor`):
+ Input tensor, of which the average must be computed.
+ weights (`torch.FloatTensor`, *optional*):
+ Weights tensor, of the same shape as `input_tensor`.
+ dim (`int`, *optional*):
+ The dim along which to average `input_tensor`.
+
+ Returns:
+ `torch.FloatTensor`: The tensor with values averaged along the specified `dim`.
+ """
+ if weights is not None:
+ weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))
+ sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)
+ return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights
+ else:
+ return input_tensor.mean(dim=dim)
+
+
+class PatchTSMixerForPrediction(PatchTSMixerPreTrainedModel):
+ r"""
+ `PatchTSMixer` for forecasting application.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ Returns:
+ `None`.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__(config)
+ self.loss = config.loss
+ self.use_return_dict = config.use_return_dict
+ self.prediction_channel_indices = config.prediction_channel_indices
+ self.num_parallel_samples = config.num_parallel_samples
+
+ if config.loss == "mse":
+ self.distribution_output = None
+ else:
+ dim = config.prediction_length
+ distribution_output_map = {
+ "student_t": StudentTOutput,
+ "normal": NormalOutput,
+ "negative_binomial": NegativeBinomialOutput,
+ }
+ output_class = distribution_output_map.get(config.distribution_output, None)
+ if output_class is not None:
+ self.distribution_output = output_class(dim=dim)
+ else:
+ raise ValueError(f"Unknown distribution output {config.distribution_output}")
+
+ self.model = PatchTSMixerModel(config)
+ self.head = PatchTSMixerForPredictionHead(
+ config=config,
+ distribution_output=self.distribution_output,
+ )
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(PATCHTSMIXER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=PatchTSMixerForPredictionOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ observed_mask: Optional[torch.Tensor] = None,
+ future_values: Optional[torch.Tensor] = None,
+ output_hidden_states: Optional[bool] = False,
+ return_loss: bool = True,
+ return_dict: Optional[bool] = None,
+ ) -> PatchTSMixerForPredictionOutput:
+ r"""
+ observed_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+ future_values (`torch.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting,:
+ `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target
+ values of the time series, that serve as labels for the model. The `future_values` is what the
+ Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT
+ required for a pretraining task.
+
+ For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want
+ to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter,
+ pass the target data with all channels, as channel Filtering for both prediction and target will be
+ manually applied before the loss computation.
+ return_loss (`bool`, *optional*):
+ Whether to return the loss in the `forward` call.
+
+ Returns:
+
+ """
+ if self.loss == "mse":
+ loss = nn.MSELoss(reduction="mean")
+ elif self.loss == "nll":
+ loss = nll
+ else:
+ raise ValueError("Invalid loss function: Allowed values: mse and nll")
+
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+
+ # past_values: tensor [batch_size x context_length x num_input_channels]
+ model_output = self.model(
+ past_values,
+ observed_mask=observed_mask,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ ) # model_output: [batch_size x nvars x num_patch x d_model]
+ if isinstance(model_output, tuple):
+ model_output = PatchTSMixerModelOutput(*model_output)
+
+ # tensor [batch_size x prediction_length x num_input_channels]
+ y_hat = self.head(model_output.last_hidden_state)
+
+ loss_val = None
+ if self.prediction_channel_indices is not None:
+ if self.distribution_output:
+ distribution = self.distribution_output.distribution(
+ y_hat,
+ loc=model_output.loc[..., self.prediction_channel_indices],
+ scale=model_output.scale[..., self.prediction_channel_indices],
+ )
+ if future_values is not None and return_loss is True:
+ loss_val = loss(
+ distribution,
+ future_values[..., self.prediction_channel_indices],
+ )
+ # take average of the loss
+ loss_val = weighted_average(loss_val)
+ else:
+ y_hat = (
+ y_hat * model_output.scale[..., self.prediction_channel_indices]
+ + model_output.loc[..., self.prediction_channel_indices]
+ )
+ if future_values is not None and return_loss is True:
+ loss_val = loss(y_hat, future_values[..., self.prediction_channel_indices])
+ else:
+ if self.distribution_output:
+ distribution = self.distribution_output.distribution(
+ y_hat, loc=model_output.loc, scale=model_output.scale
+ )
+ if future_values is not None and return_loss is True:
+ loss_val = loss(distribution, future_values)
+ loss_val = weighted_average(loss_val)
+ else:
+ y_hat = y_hat * model_output.scale + model_output.loc
+ if future_values is not None and return_loss is True:
+ loss_val = loss(y_hat, future_values)
+
+ if self.prediction_channel_indices is not None:
+ loc = model_output.loc[..., self.prediction_channel_indices]
+ scale = model_output.scale[..., self.prediction_channel_indices]
+ else:
+ loc = model_output.loc
+ scale = model_output.scale
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ loss_val,
+ y_hat,
+ model_output.last_hidden_state,
+ model_output.hidden_states,
+ loc,
+ scale,
+ ]
+ )
+
+ return PatchTSMixerForPredictionOutput(
+ loss=loss_val,
+ prediction_outputs=y_hat, # tensor [batch_size x prediction_length x num_input_channels]
+ last_hidden_state=model_output.last_hidden_state, # x: [batch_size x nvars x num_patch x d_model]
+ hidden_states=model_output.hidden_states,
+ loc=loc,
+ scale=scale,
+ )
+
+ def generate(
+ self,
+ past_values: torch.Tensor,
+ observed_mask: Optional[torch.Tensor] = None,
+ ) -> SamplePatchTSMixerPredictionOutput:
+ """
+ Generate sequences of sample predictions from a model with a probability distribution head.
+
+ Args:
+ past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Past values of the time series that serves as context in order to predict the future.
+
+ observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+ Return:
+ [`SamplePatchTSMixerPredictionOutput`] where the outputs `sequences` tensor will have shape `(batch_size,
+ number of samples, prediction_length, num_input_channels)`.
+ """
+ # get number of samples
+ num_parallel_samples = self.num_parallel_samples
+
+ # get model output
+ outputs = self(
+ past_values=past_values,
+ future_values=None,
+ observed_mask=observed_mask,
+ output_hidden_states=False,
+ )
+
+ # get distribution
+
+ distribution = self.distribution_output.distribution(
+ outputs.prediction_outputs, loc=outputs.loc, scale=outputs.scale
+ )
+
+ # get samples: list of [batch_size x prediction_length x num_channels]
+ samples = [distribution.sample() for _ in range(num_parallel_samples)]
+
+ # stack tensors
+ samples = torch.stack(samples, dim=1) # [batch_size x num_samples x prediction_length x num_channels]
+ return SamplePatchTSMixerPredictionOutput(sequences=samples)
+
+
+@dataclass
+class PatchTSMixerForTimeSeriesClassificationOutput(ModelOutput):
+ """
+ Output type of [`PatchTSMixerForTimeSeriesClassificationOutput`].
+
+ Args:
+ prediction_outputs (`torch.FloatTensor` of shape `(batch_size, num_labels)`):
+ Prediction output from the classfication head.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`):
+ Backbone embeddings before passing through the head.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+ loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`):
+ Total loss.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ prediction_outputs: Optional[torch.FloatTensor] = None
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class PatchTSMixerForTimeSeriesClassification(PatchTSMixerPreTrainedModel):
+ r"""
+ `PatchTSMixer` for classification application.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ Returns:
+ `None`.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__(config)
+
+ self.model = PatchTSMixerModel(config)
+ self.head = PatchTSMixerLinearHead(
+ config=config,
+ )
+ self.use_return_dict = config.use_return_dict
+ if config.scaling in ["std", "mean", True]:
+ self.inject_scale = InjectScalerStatistics4D(d_model=config.d_model, num_patches=config.num_patches)
+ else:
+ self.inject_scale = None
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(PATCHTSMIXER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(
+ output_type=PatchTSMixerForTimeSeriesClassificationOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ target_values: Optional[torch.Tensor] = None,
+ output_hidden_states: Optional[bool] = False,
+ return_loss: bool = True,
+ return_dict: Optional[bool] = None,
+ ) -> PatchTSMixerForTimeSeriesClassificationOutput:
+ r"""
+ target_values (`torch.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting,
+ `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target
+ values of the time series, that serve as labels for the model. The `target_values` is what the
+ Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT
+ required for a pretraining task.
+
+ For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want
+ to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter,
+ pass the target data with all channels, as channel Filtering for both prediction and target will be
+ manually applied before the loss computation.
+
+ For a classification task, it has a shape of `(batch_size,)`.
+
+ For a regression task, it has a shape of `(batch_size, num_targets)`.
+ return_loss (`bool`, *optional*):
+ Whether to return the loss in the `forward` call.
+
+ Returns:
+
+ """
+
+ loss = torch.nn.CrossEntropyLoss()
+
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+
+ model_output = self.model(
+ past_values,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ ) # x: [batch_size x nvars x num_patch x d_model]
+ if isinstance(model_output, tuple):
+ model_output = PatchTSMixerModelOutput(*model_output)
+
+ if self.inject_scale is not None:
+ model_output.last_hidden_state = self.inject_scale(
+ model_output.last_hidden_state,
+ loc=model_output.loc,
+ scale=model_output.scale,
+ ) # x: [batch_size x nvars x num_patch x d_model]
+
+ y_hat = self.head(model_output.last_hidden_state) # tensor [batch_size x n_labels]
+
+ if target_values is not None and return_loss is True:
+ loss_val = loss(y_hat, target_values)
+ else:
+ loss_val = None
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ loss_val,
+ y_hat,
+ model_output.last_hidden_state,
+ model_output.hidden_states,
+ ]
+ )
+
+ return PatchTSMixerForTimeSeriesClassificationOutput(
+ loss=loss_val,
+ prediction_outputs=y_hat, # tensor [batch_size x n_labels]
+ last_hidden_state=model_output.last_hidden_state, # x: [batch_size x nvars x num_patch x d_model]
+ hidden_states=model_output.hidden_states,
+ )
+
+
+@dataclass
+class PatchTSMixerForRegressionOutput(ModelOutput):
+ """
+ Output type of [`PatchTSMixerForRegressionOutput`].
+
+ Args:
+ regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`):
+ Prediction output from the regression head.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`):
+ Backbone embeddings before passing through the head.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+ loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`):
+ Total loss.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ regression_outputs: Optional[torch.FloatTensor] = None
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class InjectScalerStatistics4D(nn.Module):
+ def __init__(self, d_model: int, num_patches: int, expansion: int = 2):
+ super().__init__()
+
+ self.inverse_trans_expansion = nn.Linear(d_model + 2, expansion * d_model)
+ self.inverse_trans_compression = nn.Linear(expansion * d_model, d_model)
+ self.map_scale_expansion = nn.Linear(2, 2 * expansion)
+ self.map_scale_compression = nn.Linear(2 * expansion, 2)
+ self.num_patches = num_patches
+
+ def forward(self, inputs: torch.Tensor, loc: torch.Tensor, scale: torch.Tensor):
+ """
+ Args:
+ inputs (`torch.Tensor` of shape `(batch_size, num_input_channels, num_patch, d_model)`)
+ loc (`torch.Tensor` of shape `(batch_size, 1, num_input_channels)`)
+ scale (`torch.Tensor` of shape `(batch_size, 1, num_input_channels)`)
+ Returns:
+ `torch.Tensor` of shape `(batch_size, num_input_channels, num_patch, d_model)`
+ """
+
+ mean = loc.transpose(-1, -2) # [batch_size x n_channels x 1 ]
+ mean = mean.unsqueeze(-2) # [batch_size x n_channels x 1 x 1]
+ mean = mean.repeat(1, 1, self.num_patches, 1) # [batch_size x n_channels x num_patch x 1]
+
+ stdev = scale.transpose(-1, -2) # [batch_size x n_channels x 1 ]
+ stdev = stdev.unsqueeze(-2) # [batch_size x n_channels x 1 x 1]
+ stdev = stdev.repeat(1, 1, self.num_patches, 1) # [batch_size x n_channels x num_patch x 1]
+
+ concat_stats = torch.cat([mean, stdev], dim=-1) # [batch_size x n_channels x num_patch x 2]
+
+ concat_stats = self.map_scale_expansion(concat_stats) # [batch_size x n_channels x num_patch x (2*expansion)]
+ concat_stats = self.map_scale_compression(concat_stats) # [batch_size x n_channels x num_patch x 2]
+
+ inputs = torch.cat([inputs, concat_stats], dim=-1) # [batch_size x channels x num_patch x d_model+2]
+ inputs = self.inverse_trans_expansion(inputs) # [batch_size x channels x num_patch x (expansion*d_model)]
+ inputs = self.inverse_trans_compression(inputs) # [batch_size x channels x num_patch x d_model]
+
+ return inputs
+
+
+class PatchTSMixerForRegression(PatchTSMixerPreTrainedModel):
+ r"""
+ `PatchTSMixer` for regression application.
+
+ Args:
+ config (`PatchTSMixerConfig`):
+ Configuration.
+
+ Returns:
+ `None`.
+ """
+
+ def __init__(self, config: PatchTSMixerConfig):
+ super().__init__(config)
+
+ self.model = PatchTSMixerModel(config)
+
+ self.loss = config.loss
+ self.distribution_output = config.distribution_output
+
+ self.use_return_dict = config.use_return_dict
+ self.num_parallel_samples = config.num_parallel_samples
+
+ if config.loss == "mse":
+ self.distribution_output = None
+ else:
+ distribution_output_map = {
+ "student_t": StudentTOutput,
+ "normal": NormalOutput,
+ "negative_binomial": NegativeBinomialOutput,
+ }
+ output_class = distribution_output_map.get(config.distribution_output)
+ if output_class is not None:
+ self.distribution_output = output_class(dim=config.num_targets)
+ else:
+ raise ValueError(f"Unknown distribution output {config.distribution_output}")
+
+ if config.scaling in ["std", "mean", True]:
+ self.inject_scale = InjectScalerStatistics4D(d_model=config.d_model, num_patches=config.num_patches)
+ else:
+ self.inject_scale = None
+
+ self.head = PatchTSMixerLinearHead(
+ config=config,
+ distribution_output=self.distribution_output,
+ )
+
+ # Initialize weights and apply final processing
+ if config.post_init:
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(PATCHTSMIXER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=PatchTSMixerForRegressionOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ target_values: Optional[torch.Tensor] = None,
+ output_hidden_states: Optional[bool] = False,
+ return_loss: bool = True,
+ return_dict: Optional[bool] = None,
+ ) -> PatchTSMixerForRegressionOutput:
+ r"""
+ target_values (`torch.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting,
+ `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target
+ values of the time series, that serve as labels for the model. The `target_values` is what the
+ Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT
+ required for a pretraining task.
+
+ For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want
+ to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter,
+ pass the target data with all channels, as channel Filtering for both prediction and target will be
+ manually applied before the loss computation.
+
+ For a classification task, it has a shape of `(batch_size,)`.
+
+ For a regression task, it has a shape of `(batch_size, num_targets)`.
+ return_loss (`bool`, *optional*):
+ Whether to return the loss in the `forward` call.
+
+ Returns:
+
+ """
+
+ if self.loss == "mse":
+ loss = nn.MSELoss(reduction="mean")
+ elif self.loss == "nll":
+ loss = nll
+ else:
+ raise ValueError("Invalid loss function: Allowed values: mse and nll")
+
+ return_dict = return_dict if return_dict is not None else self.use_return_dict
+ model_output = self.model(
+ past_values,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ ) # model_output: [batch_size x nvars x num_patch x d_model]
+ if isinstance(model_output, tuple):
+ model_output = PatchTSMixerModelOutput(*model_output)
+
+ if self.inject_scale is not None:
+ model_output.last_hidden_state = self.inject_scale(
+ model_output.last_hidden_state,
+ loc=model_output.loc,
+ scale=model_output.scale,
+ ) # x: [batch_size x nvars x num_patch x d_model]
+
+ y_hat = self.head(model_output.last_hidden_state) # [batch_size x num_targets]
+
+ if target_values is not None and return_loss is True:
+ if self.distribution_output:
+ if self.distribution_output == "negative_binomial" and torch.any(target_values < 0):
+ raise Exception("target_values cannot be negative for negative_binomial distribution.")
+ distribution = self.distribution_output.distribution(y_hat)
+ # y_hat should be a 2-tuple, each with dimension [bs, num_targets]
+ y_hat = tuple([item.view(-1, self.config.num_targets) for item in y_hat])
+ loss_val = loss(distribution, target_values)
+ # take average of the loss
+ loss_val = weighted_average(loss_val)
+ else:
+ loss_val = loss(y_hat, target_values)
+ else:
+ loss_val = None
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ loss_val,
+ y_hat,
+ model_output.last_hidden_state,
+ model_output.hidden_states,
+ ]
+ )
+
+ return PatchTSMixerForRegressionOutput(
+ loss=loss_val,
+ regression_outputs=y_hat, # tensor [batch_size x num_targets]
+ last_hidden_state=model_output.last_hidden_state, # [batch_size x nvars x num_patch x d_model]
+ hidden_states=model_output.hidden_states,
+ )
+
+ def generate(
+ self,
+ past_values: torch.Tensor,
+ ) -> SamplePatchTSMixerRegressionOutput:
+ """
+ Generate sequences of sample predictions from a model with a probability distribution head.
+
+ Args:
+ past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Past values of the time series that serves as context in order to predict the target values.
+
+ Return:
+ [`SamplePatchTSMixerRegressionOutput`] where the outputs `sequences` tensor will have shape `(batch_size,
+ number of samples, num_targets)`.
+ """
+ # get number of samples
+ num_parallel_samples = self.num_parallel_samples
+
+ # get model output
+ outputs = self(
+ past_values=past_values,
+ target_values=None,
+ output_hidden_states=False,
+ )
+
+ # get distribution
+ distribution = self.distribution_output.distribution(outputs.regression_outputs)
+
+ # get samples
+ samples = [
+ distribution.sample() for _ in range(num_parallel_samples)
+ ] # samples: list of [batch_size x num_targets]
+ # stack tensors
+ # [batch_size x num_samples x num_targets]
+ samples = torch.stack(samples, dim=1).view(-1, num_parallel_samples, self.config.num_targets)
+ return SamplePatchTSMixerRegressionOutput(sequences=samples)
+
+
+__all__ = [
+ "PatchTSMixerPreTrainedModel",
+ "PatchTSMixerModel",
+ "PatchTSMixerForPretraining",
+ "PatchTSMixerForPrediction",
+ "PatchTSMixerForTimeSeriesClassification",
+ "PatchTSMixerForRegression",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/__init__.py b/lib/python3.12/site-packages/transformers/models/pegasus_x/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d362bacf491ea32dcbc89e7b8bba4d9ae1a9261
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pegasus_x/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_pegasus_x import *
+ from .modeling_pegasus_x import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..277ec32117f7eaf072584395f5da6f04fd0bfdfe
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/configuration_pegasus_x.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/configuration_pegasus_x.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7d5760813cba29883d0108a435c049de7b766d13
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/configuration_pegasus_x.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/modeling_pegasus_x.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/modeling_pegasus_x.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c650f167ee91d229059a5f341d8c06167f6bad72
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pegasus_x/__pycache__/modeling_pegasus_x.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/configuration_pegasus_x.py b/lib/python3.12/site-packages/transformers/models/pegasus_x/configuration_pegasus_x.py
new file mode 100644
index 0000000000000000000000000000000000000000..c92f5662b5992f0edbd48c9d37c9151fdf13bc03
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pegasus_x/configuration_pegasus_x.py
@@ -0,0 +1,177 @@
+# coding=utf-8
+# Copyright 2022, Google 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.
+"""PEGASUS-X model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class PegasusXConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`PegasusXModel`]. It is used to instantiate a
+ PEGASUS-X 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 PEGASUS-X
+ [google/pegasus-x-large](https://huggingface.co/google/pegasus-x-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 96103):
+ Vocabulary size of the PEGASUS-X model. Defines the number of different tokens that can be represented by
+ the `inputs_ids` passed when calling [`PegasusXModel`].
+ d_model (`int`, *optional*, defaults to 1024):
+ Dimension of the layers and the pooler layer.
+ encoder_layers (`int`, *optional*, defaults to 16):
+ Number of encoder layers.
+ decoder_layers (`int`, *optional*, defaults to 16):
+ 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):
+ Dimension of the "intermediate" (often named feed-forward) layer in decoder.
+ encoder_ffn_dim (`int`, *optional*, defaults to 4096):
+ Dimension of the "intermediate" (often named feed-forward) layer in decoder.
+ activation_function (`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.
+ 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 16384):
+ 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.
+ 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.
+ 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 1):
+ The id of the token to force as the last generated token when `max_length` is reached. Usually set to
+ `eos_token_id`.
+ num_global_tokens (`int`, *optional*, defaults to 128):
+ Number of global tokens to use for the encoder
+ block_size (`int`, *optional*, defaults to 512):
+ Block size for encoder local attention. Sequence length should be an exact multiple of block size.
+ block_size must be a multiple of 2 if stagger_local_block is True
+ stagger_local_block (`bool`, *optional*, defaults to `True`):
+ Whether to stagger every other local attention by half a block
+
+ Example:
+
+ ```python
+ >>> from transformers import PegasusXConfig, PegasusXModel
+
+ >>> # Initializing a PEGASUS google/pegasus-x-large style configuration
+ >>> configuration = PegasusXConfig()
+
+ >>> # Initializing a model (with random weights) from the google/pegasus-x-large style configuration
+ >>> model = PegasusXModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "pegasus_x"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
+
+ def __init__(
+ self,
+ vocab_size=96103,
+ max_position_embeddings=16384,
+ encoder_layers=16,
+ encoder_ffn_dim=4096,
+ encoder_attention_heads=16,
+ decoder_layers=16,
+ decoder_ffn_dim=4096,
+ decoder_attention_heads=16,
+ encoder_layerdrop=0.0,
+ decoder_layerdrop=0.0,
+ use_cache=True,
+ is_encoder_decoder=True,
+ activation_function="gelu",
+ d_model=1024,
+ dropout=0.1,
+ attention_dropout=0.0,
+ activation_dropout=0.0,
+ init_std=0.02,
+ decoder_start_token_id=0,
+ scale_embedding=True,
+ pad_token_id=0,
+ eos_token_id=1,
+ forced_eos_token_id=1,
+ num_global_tokens=32,
+ block_size=512,
+ stagger_local_blocks=True,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ 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.encoder_layerdrop = encoder_layerdrop
+ self.decoder_layerdrop = decoder_layerdrop
+ self.use_cache = use_cache
+ self.num_hidden_layers = encoder_layers
+ self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
+
+ self.num_global_tokens = num_global_tokens
+ self.block_size = block_size
+ self.stagger_local_blocks = stagger_local_blocks
+
+ super().__init__(
+ pad_token_id=pad_token_id,
+ eos_token_id=eos_token_id,
+ is_encoder_decoder=is_encoder_decoder,
+ decoder_start_token_id=decoder_start_token_id,
+ forced_eos_token_id=forced_eos_token_id,
+ **kwargs,
+ )
+
+ @property
+ def num_attention_heads(self) -> int:
+ return self.encoder_attention_heads
+
+ @property
+ def hidden_size(self) -> int:
+ return self.d_model
+
+
+__all__ = ["PegasusXConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/pegasus_x/modeling_pegasus_x.py b/lib/python3.12/site-packages/transformers/models/pegasus_x/modeling_pegasus_x.py
new file mode 100644
index 0000000000000000000000000000000000000000..646ab195947b9084bfb76b9879198fd28b2557da
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pegasus_x/modeling_pegasus_x.py
@@ -0,0 +1,1621 @@
+# coding=utf-8
+# Copyright 2022, Google 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 PEGASUS-X model."""
+
+import dataclasses
+import math
+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 ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_end_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from .configuration_pegasus_x import PegasusXConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "google/pegasus-x-base"
+_CONFIG_FOR_DOC = "PegasusXConfig"
+
+
+@dataclasses.dataclass
+class DimensionInfo:
+ """Wrapper for dimension info."""
+
+ batch_size: int # batch size
+ seq_len: int # token length
+ block_size: int # block size
+ num_heads: int # num heads
+ hidden_dim: int # hidden dim
+ dim_per_head: int # dim per head
+ num_blocks: int # num blocks
+ global_len: int # global length
+ padded_seq_len: int # padded token seq length
+
+ # Note: Compared to the original Flax implementation, we will pad the token representations to
+ # a multiple of block size at the start of the encoder layers, so T=P always.
+
+
+# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
+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()
+ 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
+
+
+# Copied from transformers.models.bart.modeling_bart.BartScaledWordEmbedding with Bart->PegasusX
+class PegasusXScaledWordEmbedding(nn.Embedding):
+ """
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+class PegasusXSinusoidalPositionalEmbedding(nn.Module):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, embed_dim, max_scale: int = 10000.0):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.max_scale = max_scale
+
+ @torch.no_grad()
+ def forward(self, input_embeds: torch.Tensor, past_key_values_length: int = 0) -> torch.Tensor:
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+ batch_size, seq_len = input_embeds.shape[:2]
+ positions = torch.arange(
+ past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=input_embeds.device
+ )[:, None]
+ pe = torch.zeros((seq_len, self.embed_dim), device=input_embeds.device, dtype=input_embeds.dtype)
+ half_d_feature = self.embed_dim // 2
+ div_term = torch.exp(
+ torch.arange(half_d_feature, device=input_embeds.device, dtype=torch.int64).type_as(input_embeds)
+ * -(np.log(float(self.max_scale)) / (half_d_feature - 1))
+ )
+ pe[:, :half_d_feature] = torch.sin(positions * div_term)
+ pe[:, half_d_feature:] = torch.cos(positions * div_term)
+ return pe[None].expand(batch_size, -1, -1)
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->PegasusX
+class PegasusXAttention(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[PegasusXConfig] = 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 PegasusXGlobalLocalAttention(nn.Module):
+ """Global + Local attention. For use with Encoder only."""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ block_size: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.block_size = block_size
+ 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}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+
+ 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,
+ token_hidden_states: torch.Tensor,
+ global_hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ output_attentions: bool = False,
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ """Input shape: Batch x Time x Channel"""
+ dim = DimensionInfo(
+ batch_size=token_hidden_states.shape[0],
+ seq_len=token_hidden_states.shape[1],
+ block_size=self.block_size,
+ num_heads=self.num_heads,
+ hidden_dim=token_hidden_states.shape[2],
+ dim_per_head=self.head_dim,
+ num_blocks=token_hidden_states.shape[1] // self.block_size,
+ global_len=global_hidden_states.shape[1],
+ padded_seq_len=token_hidden_states.shape[1],
+ )
+
+ # [batch_size, num_heads, padded_seq_len, dim_per_head]
+ local_q = self._shape(
+ self.q_proj(token_hidden_states) * self.scaling,
+ seq_len=dim.padded_seq_len,
+ bsz=dim.batch_size,
+ )
+ local_k = self._shape(
+ self.k_proj(token_hidden_states),
+ seq_len=dim.padded_seq_len,
+ bsz=dim.batch_size,
+ )
+ local_v = self._shape(
+ self.v_proj(token_hidden_states),
+ seq_len=dim.padded_seq_len,
+ bsz=dim.batch_size,
+ )
+
+ # [batch_size, num_heads, global_len, dim_per_head]
+ global_q = self._shape(
+ self.q_proj(global_hidden_states) * self.scaling,
+ seq_len=dim.global_len,
+ bsz=dim.batch_size,
+ )
+ global_k = self._shape(
+ self.k_proj(global_hidden_states),
+ seq_len=dim.global_len,
+ bsz=dim.batch_size,
+ )
+ global_v = self._shape(
+ self.v_proj(global_hidden_states),
+ seq_len=dim.global_len,
+ bsz=dim.batch_size,
+ )
+
+ global_attn_output, global_attn_probs = self.compute_global_attention_representations(
+ global_q=global_q,
+ global_k=global_k,
+ global_v=global_v,
+ local_k=local_k,
+ local_v=local_v,
+ mask=attention_mask,
+ dim=dim,
+ )
+ local_attn_output, local_attn_probs = self.compute_local_attention_representations(
+ global_k=global_k,
+ global_v=global_v,
+ local_q=local_q,
+ local_k=local_k,
+ local_v=local_v,
+ mask=attention_mask,
+ dim=dim,
+ )
+
+ # [batch_size, global_len, hidden_dim]
+ global_attn_output = (
+ global_attn_output.transpose(1, 2).contiguous().view(dim.batch_size, dim.global_len, dim.hidden_dim)
+ )
+ # [batch_size, global_len, hidden_dim]
+ global_attn_output = self.out_proj(global_attn_output)
+ # [batch_size, num_heads, block_size, num_heads, dim_per_head]
+ local_attn_output = local_attn_output.permute(0, 2, 3, 1, 4).contiguous()
+ # [batch_size, padded_seq_len, hidden_dim]
+ local_attn_output = local_attn_output.view(dim.batch_size, dim.padded_seq_len, dim.hidden_dim)
+ # [batch_size, padded_seq_len, hidden_dim]
+ local_attn_output = self.out_proj(local_attn_output)
+
+ if output_attentions:
+ attn_probs = {"global": global_attn_probs, "local": local_attn_probs}
+ else:
+ attn_probs = None
+
+ return local_attn_output, global_attn_output, attn_probs
+
+ def compute_global_attention_representations(
+ self, global_q, global_k, global_v, local_k, local_v, mask, dim: DimensionInfo
+ ):
+ """Compute attention representations for global tokens.
+
+ Global tokens will attend to both global tokens as well as all input sequence tokens. Because the input
+ sequence tokens are arranged in blocks for local attention, we unblock them and compute attention.
+
+ Args:
+ global_q (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
+ query vectors from global tokens
+ global_k (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
+ key vectors from global tokens
+ global_v (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
+ value vectors from global tokens
+ local_k (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
+ key vectors from local tokens
+ local_v (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
+ value vectors from local tokens
+ mask (`torch.FloatTensor`) of shape [batch_size, padded_seq_len]: attention mask
+ dim (DimensionInfo): DimensionInfo wrapper for dimensions
+
+ Returns:
+ output of shape `[batch_sizes, length, features]`. where length will be padded to a multiple of block_size
+ """
+ # [batch_size, num_heads, global_len+padded_seq_len, dim_per_head]
+ global_and_local_k = torch.cat([global_k, local_k], dim=2)
+ # [batch_size, num_heads, global_len+padded_seq_len, dim_per_head]
+ global_and_local_v = torch.cat([global_v, local_v], dim=2)
+
+ # [batch_size, global_len+padded_seq_len]
+ extended_mask = nn.functional.pad(mask, pad=(dim.global_len, 0), value=0)
+
+ # [batch_size, num_heads, global_len, global_len+padded_seq_len]
+ attn_weights = torch.einsum("BHGF,BHXF->BHGX", global_q, global_and_local_k)
+ attn_weights = attn_weights + extended_mask[:, None, None, :]
+ attn_probs = nn.functional.softmax(attn_weights, dim=-1)
+ attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training)
+
+ # [batch_size, num_heads, global_len, F]
+ attn_output = torch.einsum("BHGX,BHXF->BHGF", attn_probs, global_and_local_v)
+ return attn_output, attn_probs
+
+ def compute_local_attention_representations(
+ self, global_k, global_v, local_q, local_k, local_v, mask, dim: DimensionInfo
+ ):
+ """Compute attention representations for local tokens.
+
+ Local tokens will attend to both global tokens as well as all other tokens within the same local block. Hence,
+ we need to tile and concatenate the global tokens to every local block
+
+ Args:
+ global_k (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
+ key vectors from global tokens
+ global_v (`torch.FloatTensor`) of shape [batch_size, num_heads, global_len, dim_per_head]:
+ value vectors from global tokens
+ local_q (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
+ query vectors from local tokens
+ local_k (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
+ key vectors from local tokens
+ local_v (`torch.FloatTensor`) of shape [batch_size, num_heads, padded_seq_len, dim_per_head]:
+ value vectors from local tokens
+ mask (`torch.FloatTensor`) of shape [batch_size, padded_seq_len]: attention mask
+ dim (DimensionInfo): DimensionInfo wrapper for dimensions
+
+ Returns:
+ output of shape `[batch_sizes, length, features]`. where length will be padded to a multiple of block_size
+ """
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ blocked_local_q = local_q.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ blocked_local_k = local_k.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ blocked_local_v = local_v.view(dim.batch_size, dim.num_heads, dim.num_blocks, dim.block_size, dim.dim_per_head)
+
+ # [batch_size, num_blocks, global_len+block_size]
+ extended_mask = nn.functional.pad(
+ mask.view(dim.batch_size, dim.num_blocks, dim.block_size),
+ pad=(dim.global_len, 0),
+ value=0,
+ )
+
+ # [batch_size, num_heads, num_blocks, block_size, global_len]
+ blocked_local2global = torch.einsum("BHNKF,BHGF->BHNKG", blocked_local_q, global_k)
+ # [batch_size, num_heads, num_blocks, block_size, block_size]
+ blocked_local2local = torch.einsum("BHNKF,BHNXF->BHNKX", blocked_local_q, blocked_local_k)
+
+ # [batch_size, num_heads, num_blocks, block_size, global_len+block_size]
+ attn_weights = torch.cat([blocked_local2global, blocked_local2local], dim=-1)
+ attn_weights = attn_weights + extended_mask[:, None, :, None, :]
+ attn_probs = nn.functional.softmax(attn_weights, dim=-1)
+ attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training)
+
+ # [batch_size, num_heads, num_blocks, block_size, global_len]
+ local2global_attn_probs = attn_probs[:, :, :, :, : dim.global_len]
+ # [batch_size, num_heads, num_blocks, block_size, block_size]
+ local2local_attn_probs = attn_probs[:, :, :, :, dim.global_len :]
+
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ local2global_attn_output = torch.einsum("BHNKG,BHGF->BHNKF", local2global_attn_probs, global_v)
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ local2local_attn_output = torch.einsum("BHNKX,BHNXF->BHNKF", local2local_attn_probs, blocked_local_v)
+ # [batch_size, num_heads, num_blocks, block_size, dim_per_head]
+ attn_output = local2global_attn_output + local2local_attn_output
+ return attn_output, attn_probs
+
+
+class PegasusXEncoderLayer(nn.Module):
+ def __init__(self, stagger_blocks_this_layer: bool, config: PegasusXConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.self_attn = PegasusXGlobalLocalAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ block_size=config.block_size,
+ dropout=config.attention_dropout,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.global_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)
+ self.stagger_blocks_this_layer = stagger_blocks_this_layer
+ self.block_size = config.block_size
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ global_hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ output_attentions: bool = False,
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape *(seq_len, batch, embed_dim)*
+ global_hidden_states (`torch.FloatTensor`): global token hidden states
+ *(seq_len, num_global_tokens, 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.
+ 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
+ global_residual = global_hidden_states
+
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ global_hidden_states = self.global_self_attn_layer_norm(global_hidden_states)
+
+ if self.stagger_blocks_this_layer:
+ # Pad the blocks to simulate staggering
+ hidden_states, attention_mask = self.pad_local_tokens(
+ hidden_states=hidden_states, attention_mask=attention_mask, block_size=self.block_size
+ )
+
+ hidden_states, global_hidden_states, attn_weights = self.self_attn(
+ token_hidden_states=hidden_states,
+ global_hidden_states=global_hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ )
+
+ if self.stagger_blocks_this_layer:
+ # Undo the padding
+ hidden_states = self.unpad_local_tokens(padded_hidden_states=hidden_states, block_size=self.block_size)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training)
+ global_hidden_states = global_residual + global_hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(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
+
+ global_residual = global_hidden_states
+ global_hidden_states = self.final_layer_norm(global_hidden_states)
+ global_hidden_states = self.activation_fn(self.fc1(global_hidden_states))
+ global_hidden_states = nn.functional.dropout(
+ global_hidden_states, p=self.activation_dropout, training=self.training
+ )
+ global_hidden_states = self.fc2(global_hidden_states)
+ global_hidden_states = nn.functional.dropout(global_hidden_states, p=self.dropout, training=self.training)
+ global_hidden_states = global_residual + global_hidden_states
+ outputs = (hidden_states, global_hidden_states)
+
+ if output_attentions:
+ outputs += (attn_weights,)
+
+ return outputs
+
+ @classmethod
+ def pad_local_tokens(cls, hidden_states, attention_mask, block_size):
+ # hidden_states: [batch_size, seq_len, hidden_dim]
+ pad_size = block_size // 2
+ mask_min_value = torch.finfo(hidden_states.dtype).min
+ padded_hidden_states = torch.nn.functional.pad(
+ hidden_states,
+ pad=(0, 0, pad_size, pad_size),
+ )
+ padded_mask = torch.nn.functional.pad(
+ attention_mask,
+ pad=(pad_size, pad_size),
+ value=mask_min_value,
+ )
+ return padded_hidden_states, padded_mask
+
+ @classmethod
+ def unpad_local_tokens(cls, padded_hidden_states, block_size):
+ # padded_hidden_states: [batch_size, padded seq_len, hidden_dim]
+ pad_size = block_size // 2
+ return padded_hidden_states[:, pad_size:-pad_size, :]
+
+
+class PegasusXDecoderLayer(nn.Module):
+ def __init__(self, config: PegasusXConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = PegasusXAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ )
+ 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 = PegasusXAttention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ )
+ 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,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = True,
+ ) -> torch.Tensor:
+ """
+ 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, tgt_len, src_len)* where padding elements are indicated by very large negative values.
+ 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, tgt_len, src_len)* where padding elements are indicated by very large negative values.
+ past_key_value (`Tuple(torch.FloatTensor)`): 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.
+ use_cache: Whether to us KV cache for decoding
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # 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
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_value=self_attn_past_key_value,
+ 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
+
+ # Cross-Attention Block
+ cross_attn_present_key_value = None
+ cross_attn_weights = None
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
+ cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+ hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_value=cross_attn_past_key_value,
+ output_attentions=output_attentions,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # add cross-attn to positions 3,4 of present_key_value tuple
+ present_key_value = present_key_value + cross_attn_present_key_value
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(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
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights, cross_attn_weights)
+
+ if use_cache:
+ outputs += (present_key_value,)
+
+ return outputs
+
+
+class PegasusXPreTrainedModel(PreTrainedModel):
+ config_class = PegasusXConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = [r"PegasusXEncoderLayer", r"PegasusXDecoderLayer"]
+
+ 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, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+
+
+PEGASUS_X_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 ([`PegasusXConfig`]):
+ 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.
+"""
+
+PEGASUS_X_GENERATION_EXAMPLE = r"""
+ Summarization example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, PegasusXForConditionalGeneration
+
+ >>> model = PegasusXForConditionalGeneration.from_pretrained("google/pegasus-x-base")
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/pegasus-x-large")
+
+ >>> ARTICLE_TO_SUMMARIZE = (
+ ... "PG&E stated it scheduled the blackouts in response to forecasts for high winds "
+ ... "amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were "
+ ... "scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."
+ ... )
+ >>> inputs = tokenizer(ARTICLE_TO_SUMMARIZE, max_length=1024, return_tensors="pt")
+
+ >>> # Generate Summary
+ >>> summary_ids = model.generate(inputs["input_ids"])
+ >>> tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "California's largest electricity provider has turned off power to hundreds of thousands of customers."
+ ```
+"""
+
+PEGASUS_X_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)
+ 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.
+ 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)
+
+ PEGASUS-X uses the `pad_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.LongTensor` 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(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.
+ 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.
+ 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*):
+ 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 PegasusXEncoder(PegasusXPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`PegasusXEncoderLayer`].
+
+ Args:
+ config: PegasusXConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ def __init__(self, config: PegasusXConfig, embed_tokens: Optional[nn.Embedding] = None):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+
+ embed_dim = config.d_model
+ padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ if embed_tokens is not None:
+ self.embed_tokens = embed_tokens
+ else:
+ self.embed_tokens = PegasusXScaledWordEmbedding(
+ config.vocab_size, embed_dim, padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_global = nn.Embedding(config.num_global_tokens, embed_dim)
+ self.embed_positions = PegasusXSinusoidalPositionalEmbedding(embed_dim)
+ self.layers = nn.ModuleList(
+ [
+ PegasusXEncoderLayer(
+ stagger_blocks_this_layer=i % 2 == 1 and config.stagger_local_blocks, config=config
+ )
+ for i in range(config.encoder_layers)
+ ]
+ )
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def resize_position_embeddings(self, new_num_position_embeddings: int):
+ """
+ Resizes position embeddings matrix 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.
+ """
+ logger.info(f"Setting `config.max_position_embeddings={new_num_position_embeddings}`...")
+ self.config.max_position_embeddings = new_num_position_embeddings
+
+ self.embed_positions = PegasusXSinusoidalPositionalEmbedding(self.config.d_model)
+ self.embed_positions.to(self.device)
+
+ def get_position_embeddings(self) -> nn.Embedding:
+ """
+ Returns the position embeddings matrix
+ """
+ return self.embed_positions
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ 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)
+
+ 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.
+ """
+ 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
+
+ # retrieve input_ids and inputs_embeds
+ 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")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ embed_pos = self.embed_positions(inputs_embeds)
+
+ hidden_states = inputs_embeds + embed_pos
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ batch_size, seq_len, _ = hidden_states.shape
+
+ # Setup mask
+ if attention_mask is None:
+ attention_mask = torch.ones(*input_shape, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+ attention_mask = attention_mask.to(dtype=hidden_states.dtype)
+ mask_min_value = torch.finfo(hidden_states.dtype).min
+ inverted_mask = 1.0 - attention_mask
+ attention_mask = inverted_mask.masked_fill(
+ inverted_mask.to(torch.bool),
+ mask_min_value,
+ )
+
+ # padding to block_size
+ if seq_len % self.config.block_size != 0:
+ pad_len = self.config.block_size - seq_len % self.config.block_size
+ hidden_states = nn.functional.pad(hidden_states, pad=(0, 0, 0, pad_len), value=0)
+ attention_mask = nn.functional.pad(attention_mask, pad=(0, pad_len), value=mask_min_value)
+
+ # Global tokens
+ global_hidden_states = self.embed_global(
+ torch.arange(self.config.num_global_tokens, device=hidden_states.device)[None].expand(batch_size, -1)
+ )
+
+ encoder_states = () if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+
+ 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)
+ 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:
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ encoder_layer.__call__,
+ hidden_states,
+ global_hidden_states,
+ attention_mask,
+ output_attentions,
+ )
+ else:
+ layer_outputs = encoder_layer(
+ hidden_states,
+ global_hidden_states,
+ attention_mask,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+ global_hidden_states = layer_outputs[1]
+
+ if output_attentions:
+ all_attentions = all_attentions + (layer_outputs[2],)
+
+ # Undo padding-to-block-size
+ hidden_states = hidden_states[:, :seq_len]
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ if output_hidden_states:
+ encoder_states = encoder_states + ((hidden_states, global_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 PegasusXDecoder(PegasusXPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`PegasusDecoderLayer`]
+
+ Args:
+ config: PegasusXConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ def __init__(self, config: PegasusXConfig, embed_tokens: Optional[nn.Embedding] = None):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ self.max_target_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+ padding_idx = config.pad_token_id
+
+ if embed_tokens is not None:
+ self.embed_tokens = embed_tokens
+ else:
+ self.embed_tokens = PegasusXScaledWordEmbedding(
+ config.vocab_size, config.d_model, padding_idx=padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = PegasusXSinusoidalPositionalEmbedding(config.d_model)
+ self.layers = nn.ModuleList([PegasusXDecoderLayer(config) for _ in range(config.decoder_layers)])
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ 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
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ 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)
+ 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 tokens indices of encoder input_ids. 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)
+
+ 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.
+ 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
+ )
+ 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:
+ 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 decoder_input_ids or decoder_inputs_embeds")
+
+ # 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 inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask, input_shape, inputs_embeds, past_key_values_length
+ )
+
+ # expand encoder attention mask
+ if encoder_hidden_states is not None and encoder_attention_mask is not None:
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ encoder_attention_mask = _prepare_4d_attention_mask(
+ encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+ )
+
+ # embed positions
+ positions = self.embed_positions(inputs_embeds, past_key_values_length)
+
+ positions = positions.to(inputs_embeds.device)
+
+ hidden_states = inputs_embeds + positions
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ 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
+
+ # 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
+ next_decoder_cache = () if use_cache 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
+
+ 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,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ None,
+ output_attentions,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ 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[3 if output_attentions else 1],)
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if encoder_hidden_states is not None:
+ all_cross_attentions += (layer_outputs[2],)
+
+ hidden_states = self.layer_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, all_cross_attentions]
+ if v is not None
+ )
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+@add_start_docstrings(
+ "The bare PEGASUS-X Model outputting raw hidden-states without any specific head on top.",
+ PEGASUS_X_START_DOCSTRING,
+)
+class PegasusXModel(PegasusXPreTrainedModel):
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
+
+ def __init__(self, config: PegasusXConfig):
+ super().__init__(config)
+
+ vocab_size = config.vocab_size
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+ padding_idx = config.pad_token_id
+ self.shared = PegasusXScaledWordEmbedding(
+ vocab_size, config.d_model, padding_idx=padding_idx, embed_scale=embed_scale
+ )
+
+ self.encoder = PegasusXEncoder(config, self.shared)
+ self.decoder = PegasusXDecoder(config, self.shared)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.shared
+
+ def set_input_embeddings(self, value):
+ self.shared = value
+ self.encoder.embed_tokens = self.shared
+ self.decoder.embed_tokens = self.shared
+
+ def get_encoder(self):
+ return self.encoder
+
+ def get_decoder(self):
+ return self.decoder
+
+ def resize_position_embeddings(self, new_num_position_embeddings: int):
+ """
+ Resizes position embeddings matrix 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.config.max_position_embeddings = new_num_position_embeddings
+ self.encoder.resize_position_embeddings(new_num_position_embeddings)
+ self.decoder.resize_position_embeddings(new_num_position_embeddings)
+
+ def get_position_embeddings(self) -> Tuple[nn.Embedding]:
+ """
+ Returns the position embeddings matrix
+ """
+ return (self.encoder.get_position_embeddings(), self.decoder.get_position_embeddings())
+
+ @add_start_docstrings_to_model_forward(PEGASUS_X_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ decoder_input_ids: Optional[torch.Tensor] = None,
+ decoder_attention_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,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, Seq2SeqModelOutput]:
+ r"""
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, PegasusModel
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/pegasus-x-large")
+ >>> model = PegasusModel.from_pretrained("google/pegasus-x-large")
+
+ >>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt")
+ >>> decoder_inputs = tokenizer("Studies show that", return_tensors="pt")
+ >>> outputs = model(input_ids=inputs.input_ids, decoder_input_ids=decoder_inputs.input_ids)
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 4, 1024]
+ ```"""
+
+ 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 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,
+ )
+ # 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,
+ )
+
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
+ decoder_outputs = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs[0],
+ encoder_attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ 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,
+ )
+
+
+@add_start_docstrings("The PEGASUS-X for conditional generation (e.g. summarization).", PEGASUS_X_START_DOCSTRING)
+class PegasusXForConditionalGeneration(PegasusXPreTrainedModel, GenerationMixin):
+ base_model_prefix = "model"
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]
+
+ def __init__(self, config: PegasusXConfig):
+ super().__init__(config)
+ self.model = PegasusXModel(config)
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_encoder(self):
+ return self.model.get_encoder()
+
+ def get_decoder(self):
+ return self.model.get_decoder()
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+ def resize_position_embeddings(self, new_num_position_embeddings: int):
+ """
+ Resizes position embeddings matrix 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.config.max_position_embeddings = new_num_position_embeddings
+ self.model.encoder.resize_position_embeddings(new_num_position_embeddings)
+ self.model.decoder.resize_position_embeddings(new_num_position_embeddings)
+
+ def get_position_embeddings(self) -> Tuple[nn.Embedding]:
+ """
+ Returns the position embeddings matrix
+ """
+ return (self.model.encoder.get_position_embeddings(), self.model.decoder.get_position_embeddings())
+
+ @add_start_docstrings_to_model_forward(PEGASUS_X_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
+ @add_end_docstrings(PEGASUS_X_GENERATION_EXAMPLE)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ decoder_input_ids: Optional[torch.Tensor] = None,
+ decoder_attention_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.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, 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:
+ if use_cache:
+ logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
+ use_cache = False
+ if 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
+ )
+
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ lm_logits = self.lm_head(outputs[0])
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.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_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)
+
+ @staticmethod
+ def _reorder_cache(past_key_values, beam_idx):
+ reordered_past = ()
+ for layer_past in past_key_values:
+ # cached cross_attention states don't have to be reordered -> they are always the same
+ 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
+
+
+# Copied from transformers.models.bart.modeling_bart.BartDecoderWrapper with Bart->PegasusX
+class PegasusXDecoderWrapper(PegasusXPreTrainedModel):
+ """
+ This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is
+ used in combination with the [`EncoderDecoderModel`] framework.
+ """
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.decoder = PegasusXDecoder(config)
+
+ def forward(self, *args, **kwargs):
+ return self.decoder(*args, **kwargs)
+
+
+__all__ = ["PegasusXForConditionalGeneration", "PegasusXModel", "PegasusXPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__init__.py b/lib/python3.12/site-packages/transformers/models/poolformer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0d79274910e3480302eef831f2524cc04f40d86
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/poolformer/__init__.py
@@ -0,0 +1,29 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_poolformer import *
+ from .feature_extraction_poolformer import *
+ from .image_processing_poolformer import *
+ from .modeling_poolformer import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fa82774343cb99d86cf278674f576a0816b5babc
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/configuration_poolformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/configuration_poolformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9611e3806c56a3770d582ba2952e01472b278673
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/configuration_poolformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/feature_extraction_poolformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/feature_extraction_poolformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..71b90b425a198ef9923d92aee569878407c46e4d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/feature_extraction_poolformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/image_processing_poolformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/image_processing_poolformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bfd294b5954e4b8894f8f29869d7c8cd25b875c5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/image_processing_poolformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/modeling_poolformer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/modeling_poolformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f660b5c3190283100e71abb7536e9e3cc9e1a5f7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/poolformer/__pycache__/modeling_poolformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/configuration_poolformer.py b/lib/python3.12/site-packages/transformers/models/poolformer/configuration_poolformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..cdaf1306313830982fd799d9bdc24dd4f1a7412a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/poolformer/configuration_poolformer.py
@@ -0,0 +1,148 @@
+# coding=utf-8
+# Copyright 2022 Sea AI 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.
+"""PoolFormer 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__)
+
+
+class PoolFormerConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of [`PoolFormerModel`]. It is used to instantiate a
+ PoolFormer 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 PoolFormer
+ [sail/poolformer_s12](https://huggingface.co/sail/poolformer_s12) 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 channels in the input image.
+ patch_size (`int`, *optional*, defaults to 16):
+ The size of the input patch.
+ stride (`int`, *optional*, defaults to 16):
+ The stride of the input patch.
+ pool_size (`int`, *optional*, defaults to 3):
+ The size of the pooling window.
+ mlp_ratio (`float`, *optional*, defaults to 4.0):
+ The ratio of the number of channels in the output of the MLP to the number of channels in the input.
+ depths (`list`, *optional*, defaults to `[2, 2, 6, 2]`):
+ The depth of each encoder block.
+ hidden_sizes (`list`, *optional*, defaults to `[64, 128, 320, 512]`):
+ The hidden sizes of each encoder block.
+ patch_sizes (`list`, *optional*, defaults to `[7, 3, 3, 3]`):
+ The size of the input patch for each encoder block.
+ strides (`list`, *optional*, defaults to `[4, 2, 2, 2]`):
+ The stride of the input patch for each encoder block.
+ padding (`list`, *optional*, defaults to `[2, 1, 1, 1]`):
+ The padding of the input patch for each encoder block.
+ num_encoder_blocks (`int`, *optional*, defaults to 4):
+ The number of encoder blocks.
+ drop_path_rate (`float`, *optional*, defaults to 0.0):
+ The dropout rate for the dropout layers.
+ hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The activation function for the hidden layers.
+ use_layer_scale (`bool`, *optional*, defaults to `True`):
+ Whether to use layer scale.
+ layer_scale_init_value (`float`, *optional*, defaults to 1e-05):
+ The initial value for the layer scale.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The initializer range for the weights.
+
+ Example:
+
+ ```python
+ >>> from transformers import PoolFormerConfig, PoolFormerModel
+
+ >>> # Initializing a PoolFormer sail/poolformer_s12 style configuration
+ >>> configuration = PoolFormerConfig()
+
+ >>> # Initializing a model (with random weights) from the sail/poolformer_s12 style configuration
+ >>> model = PoolFormerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "poolformer"
+
+ def __init__(
+ self,
+ num_channels=3,
+ patch_size=16,
+ stride=16,
+ pool_size=3,
+ mlp_ratio=4.0,
+ depths=[2, 2, 6, 2],
+ hidden_sizes=[64, 128, 320, 512],
+ patch_sizes=[7, 3, 3, 3],
+ strides=[4, 2, 2, 2],
+ padding=[2, 1, 1, 1],
+ num_encoder_blocks=4,
+ drop_path_rate=0.0,
+ hidden_act="gelu",
+ use_layer_scale=True,
+ layer_scale_init_value=1e-5,
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ self.num_channels = num_channels
+ self.patch_size = patch_size
+ self.stride = stride
+ self.padding = padding
+ self.pool_size = pool_size
+ self.hidden_sizes = hidden_sizes
+ self.mlp_ratio = mlp_ratio
+ self.depths = depths
+ self.patch_sizes = patch_sizes
+ self.strides = strides
+ self.num_encoder_blocks = num_encoder_blocks
+ self.drop_path_rate = drop_path_rate
+ self.hidden_act = hidden_act
+ self.use_layer_scale = use_layer_scale
+ self.layer_scale_init_value = layer_scale_init_value
+ self.initializer_range = initializer_range
+ super().__init__(**kwargs)
+
+
+class PoolFormerOnnxConfig(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 2e-3
+
+
+__all__ = ["PoolFormerConfig", "PoolFormerOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/feature_extraction_poolformer.py b/lib/python3.12/site-packages/transformers/models/poolformer/feature_extraction_poolformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab4337f91f4f0a37ed687385d3c8a252b1d61116
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/poolformer/feature_extraction_poolformer.py
@@ -0,0 +1,36 @@
+# 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 PoolFormer."""
+
+import warnings
+
+from ...utils import logging
+from .image_processing_poolformer import PoolFormerImageProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+class PoolFormerFeatureExtractor(PoolFormerImageProcessor):
+ def __init__(self, *args, **kwargs) -> None:
+ warnings.warn(
+ "The class PoolFormerFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
+ " Please use PoolFormerImageProcessor instead.",
+ FutureWarning,
+ )
+ super().__init__(*args, **kwargs)
+
+
+__all__ = ["PoolFormerFeatureExtractor"]
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/image_processing_poolformer.py b/lib/python3.12/site-packages/transformers/models/poolformer/image_processing_poolformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd4c4bb77086a53ef7583d2569c24ff62c79cb20
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/poolformer/image_processing_poolformer.py
@@ -0,0 +1,360 @@
+# 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 PoolFormer."""
+
+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_scaled_image,
+ make_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
+
+
+if is_vision_available():
+ import PIL
+
+
+logger = logging.get_logger(__name__)
+
+
+class PoolFormerImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a PoolFormer 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. Can be overridden by `size` in the `preprocess` method. If crop_pct is
+ unset:
+ - size is `{"height": h, "width": w}`: the image is resized to `(h, w)`.
+ - size is `{"shortest_edge": s}`: the shortest edge of the image is resized to s whilst maintaining the
+ aspect ratio.
+
+ If crop_pct is set:
+ - size is `{"height": h, "width": w}`: the image is resized to `(int(floor(h/crop_pct)),
+ int(floor(w/crop_pct)))`
+ - size is `{"height": c, "width": c}`: the shortest edge of the image is resized to `int(floor(c/crop_pct)`
+ whilst maintaining the aspect ratio.
+ - size is `{"shortest_edge": c}`: the shortest edge of the image is resized to `int(floor(c/crop_pct)`
+ whilst maintaining the aspect ratio.
+ crop_pct (`float`, *optional*, defaults to 0.9):
+ Percentage of the image to crop from the center. Can be overridden by `crop_pct` in the `preprocess`
+ method.
+ resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+ Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method.
+ 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 the `preprocess`
+ method.
+ crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Size of the image after applying center crop. Only has an effect if `do_center_crop` is set to `True`. Can
+ be overridden by the `crop_size` 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_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`):
+ 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_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: int = 0.9,
+ resample: PILImageResampling = PILImageResampling.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 {"shortest_edge": 224}
+ size = get_size_dict(size, default_to_square=False)
+ 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.crop_pct = crop_pct
+ 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_DEFAULT_MEAN
+ self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
+
+ def resize(
+ self,
+ image: np.ndarray,
+ size: Dict[str, int],
+ crop_pct: Optional[float] = None,
+ 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.
+
+ If crop_pct is unset:
+ - size is `{"height": h, "width": w}`: the image is resized to `(h, w)`.
+ - size is `{"shortest_edge": s}`: the shortest edge of the image is resized to s whilst maintaining the
+ aspect ratio.
+
+ if crop_pct is set:
+ - size is `{"height": h, "width": w}`: the image is resized to `(int(floor(h/crop_pct)),
+ int(floor(w/crop_pct)))`
+ - size is `{"height": c, "width": c}`: the shortest edge of the image is resized to `int(floor(c/crop_pct)`
+ whilst maintaining the aspect ratio.
+ - size is `{"shortest_edge": c}`: the shortest edge of the image is resized to `int(floor(c/crop_pct)`
+ whilst maintaining the aspect ratio.
+
+ Args:
+ image (`np.ndarray`):
+ Image to resize.
+ size (`Dict[str, int]`):
+ Size of the output image.
+ crop_pct (`float`, *optional*):
+ Percentage of the image that will be cropped from the center. If set, the image is resized
+ 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 (`str` or `ChannelDimension`, *optional*):
+ The channel dimension format of the input image. If not provided, it will be inferred.
+ """
+ size = get_size_dict(size, default_to_square=False)
+ if "shortest_edge" not in size and ("height" not in size or "width" not in size):
+ raise ValueError(f"size must contain 'height' and 'width' or 'shortest_edge' as keys. Got {size.keys()}")
+ if crop_pct is not None:
+ if "shortest_edge" in size:
+ scale_size = int(size["shortest_edge"] / crop_pct)
+ elif "height" in size and "width" in size:
+ if size["height"] == size["width"]:
+ scale_size = int(size["height"] / crop_pct)
+ else:
+ scale_size = (int(size["height"] / crop_pct), int(size["width"] / crop_pct))
+ else:
+ raise ValueError("Invalid size for resize: {}".format(size))
+
+ output_size = get_resize_output_image_size(
+ image, size=scale_size, default_to_square=False, input_data_format=input_data_format
+ )
+ else:
+ if "shortest_edge" in size:
+ output_size = get_resize_output_image_size(
+ image, size=size["shortest_edge"], default_to_square=False, input_data_format=input_data_format
+ )
+ elif "height" in size and "width" in size:
+ output_size = (size["height"], size["width"])
+ else:
+ raise ValueError("Invalid size for resize: {}".format(size))
+
+ return resize(
+ image,
+ size=output_size,
+ resample=resample,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ **kwargs,
+ )
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ crop_pct: Optional[int] = None,
+ resample: PILImageResampling = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: 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,
+ return_tensors: Optional[Union[str, TensorType]] = None,
+ data_format: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> 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 applying resize.
+ crop_pct (`float`, *optional*, defaults to `self.crop_pct`):
+ Percentage of the image to crop. Only has an effect if `do_resize` is set to `True`.
+ 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 image after applying center crop.
+ 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
+ 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_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, default_to_square=False)
+ 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)
+
+ 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 do_rescale and is_scaled_image(images[0]):
+ 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_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)
+
+
+__all__ = ["PoolFormerImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/poolformer/modeling_poolformer.py b/lib/python3.12/site-packages/transformers/models/poolformer/modeling_poolformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d13304405ab2b93a63bcb5dd4618053194ba36a0
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/poolformer/modeling_poolformer.py
@@ -0,0 +1,452 @@
+# coding=utf-8
+# Copyright 2022 Sea AI Lab 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 PoolFormer model."""
+
+import collections.abc
+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, 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_poolformer import PoolFormerConfig
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "PoolFormerConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "sail/poolformer_s12"
+_EXPECTED_OUTPUT_SHAPE = [1, 512, 7, 7]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "sail/poolformer_s12"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+# 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->PoolFormer
+class PoolFormerDropPath(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 PoolFormerEmbeddings(nn.Module):
+ """
+ Construct Patch Embeddings.
+ """
+
+ def __init__(self, hidden_size, num_channels, patch_size, stride, padding, norm_layer=None):
+ super().__init__()
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ stride = stride if isinstance(stride, collections.abc.Iterable) else (stride, stride)
+ padding = padding if isinstance(padding, collections.abc.Iterable) else (padding, padding)
+
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=stride, padding=padding)
+ self.norm = norm_layer(hidden_size) if norm_layer else nn.Identity()
+
+ def forward(self, pixel_values):
+ embeddings = self.projection(pixel_values)
+ embeddings = self.norm(embeddings)
+ return embeddings
+
+
+class PoolFormerGroupNorm(nn.GroupNorm):
+ """
+ Group Normalization with 1 group. Input: tensor in shape [B, C, H, W]
+ """
+
+ def __init__(self, num_channels, **kwargs):
+ super().__init__(1, num_channels, **kwargs)
+
+
+class PoolFormerPooling(nn.Module):
+ def __init__(self, pool_size):
+ super().__init__()
+ self.pool = nn.AvgPool2d(pool_size, stride=1, padding=pool_size // 2, count_include_pad=False)
+
+ def forward(self, hidden_states):
+ return self.pool(hidden_states) - hidden_states
+
+
+class PoolFormerOutput(nn.Module):
+ def __init__(self, config, dropout_prob, hidden_size, intermediate_size):
+ super().__init__()
+ self.conv1 = nn.Conv2d(hidden_size, intermediate_size, 1)
+ self.conv2 = nn.Conv2d(intermediate_size, hidden_size, 1)
+ self.drop = PoolFormerDropPath(dropout_prob)
+ if isinstance(config.hidden_act, str):
+ self.act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.act_fn = config.hidden_act
+
+ def forward(self, hidden_states):
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.act_fn(hidden_states)
+ hidden_states = self.drop(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.drop(hidden_states)
+
+ return hidden_states
+
+
+class PoolFormerLayer(nn.Module):
+ """This corresponds to the 'PoolFormerBlock' class in the original implementation."""
+
+ def __init__(self, config, num_channels, pool_size, hidden_size, intermediate_size, drop_path):
+ super().__init__()
+ self.pooling = PoolFormerPooling(pool_size)
+ self.output = PoolFormerOutput(config, drop_path, hidden_size, intermediate_size)
+ self.before_norm = PoolFormerGroupNorm(num_channels)
+ self.after_norm = PoolFormerGroupNorm(num_channels)
+
+ # Useful for training neural nets
+ self.drop_path = PoolFormerDropPath(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((num_channels)), requires_grad=True
+ )
+ self.layer_scale_2 = nn.Parameter(
+ config.layer_scale_init_value * torch.ones((num_channels)), requires_grad=True
+ )
+
+ def forward(self, hidden_states):
+ if self.use_layer_scale:
+ pooling_output = self.pooling(self.before_norm(hidden_states))
+ scaled_op = self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * pooling_output
+ # First residual connection
+ hidden_states = hidden_states + self.drop_path(scaled_op)
+ outputs = ()
+
+ layer_output = self.output(self.after_norm(hidden_states))
+ scaled_op = self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * layer_output
+ # Second residual connection
+ output = hidden_states + self.drop_path(scaled_op)
+
+ outputs = (output,) + outputs
+ return outputs
+
+ else:
+ pooling_output = self.drop_path(self.pooling(self.before_norm(hidden_states)))
+ # First residual connection
+ hidden_states = pooling_output + hidden_states
+ outputs = ()
+
+ # Second residual connection inside the PoolFormerOutput block
+ layer_output = self.drop_path(self.output(self.after_norm(hidden_states)))
+ output = hidden_states + layer_output
+
+ outputs = (output,) + outputs
+ return outputs
+
+
+class PoolFormerEncoder(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(
+ PoolFormerEmbeddings(
+ patch_size=config.patch_sizes[i],
+ stride=config.strides[i],
+ padding=config.padding[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(
+ PoolFormerLayer(
+ config,
+ num_channels=config.hidden_sizes[i],
+ pool_size=config.pool_size,
+ hidden_size=config.hidden_sizes[i],
+ intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio),
+ drop_path=dpr[cur + j],
+ )
+ )
+ blocks.append(nn.ModuleList(layers))
+
+ self.block = nn.ModuleList(blocks)
+
+ def forward(self, pixel_values, output_hidden_states=False, return_dict=True):
+ all_hidden_states = () if output_hidden_states else None
+
+ hidden_states = pixel_values
+ for idx, layers in enumerate(zip(self.patch_embeddings, self.block)):
+ embedding_layer, block_layer = layers
+ # Get patch embeddings from hidden_states
+ hidden_states = embedding_layer(hidden_states)
+ # Send the embeddings through the blocks
+ for _, blk in enumerate(block_layer):
+ layer_outputs = blk(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 BaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
+
+
+class PoolFormerPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = PoolFormerConfig
+ base_model_prefix = "poolformer"
+ main_input_name = "pixel_values"
+ _no_split_modules = ["PoolFormerLayer"]
+
+ def _init_weights(self, 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.GroupNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+ elif isinstance(module, PoolFormerLayer):
+ if hasattr(module, "layer_scale_1"):
+ module.layer_scale_1.data.fill_(self.config.layer_scale_init_value)
+ module.layer_scale_2.data.fill_(self.config.layer_scale_init_value)
+
+
+POOLFORMER_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 ([`PoolFormerConfig`]): 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.
+"""
+
+POOLFORMER_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
+ [`PoolFormerImageProcessor.__call__`] for details.
+"""
+
+
+@add_start_docstrings(
+ "The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.",
+ POOLFORMER_START_DOCSTRING,
+)
+class PoolFormerModel(PoolFormerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.encoder = PoolFormerEncoder(config)
+
+ # 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(POOLFORMER_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=BaseModelOutputWithNoAttention,
+ config_class=_CONFIG_FOR_DOC,
+ modality="vision",
+ expected_output=_EXPECTED_OUTPUT_SHAPE,
+ )
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithNoAttention]:
+ 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, None) + encoder_outputs[1:]
+
+ return BaseModelOutputWithNoAttention(
+ last_hidden_state=sequence_output,
+ hidden_states=encoder_outputs.hidden_states,
+ )
+
+
+class PoolFormerFinalPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+
+ def forward(self, hidden_states):
+ output = self.dense(hidden_states)
+ return output
+
+
+@add_start_docstrings(
+ """
+ PoolFormer Model transformer with an image classification head on top
+ """,
+ POOLFORMER_START_DOCSTRING,
+)
+class PoolFormerForImageClassification(PoolFormerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.poolformer = PoolFormerModel(config)
+
+ # Final norm
+ self.norm = PoolFormerGroupNorm(config.hidden_sizes[-1])
+ # 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(POOLFORMER_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.poolformer(
+ pixel_values,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.classifier(self.norm(sequence_output).mean([-2, -1]))
+
+ 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)
+
+
+__all__ = ["PoolFormerForImageClassification", "PoolFormerModel", "PoolFormerPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/__init__.py b/lib/python3.12/site-packages/transformers/models/pvt/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a7448336f164f6dd226dea1dea3fff463f754c4
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pvt/__init__.py
@@ -0,0 +1,28 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_pvt import *
+ from .image_processing_pvt import *
+ from .modeling_pvt import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..454a864853896456854af9f0f3221675e072bead
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/configuration_pvt.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/configuration_pvt.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ee3372dee398ee3ec43fd35b82c6961feaf6f029
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/configuration_pvt.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/image_processing_pvt.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/image_processing_pvt.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..286f4bb6f5d49f0d205b72cb434f207ab6131053
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/image_processing_pvt.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/modeling_pvt.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/modeling_pvt.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fee68d0e38f89b60e6e292bd2faf7ae1226972a9
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/pvt/__pycache__/modeling_pvt.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/configuration_pvt.py b/lib/python3.12/site-packages/transformers/models/pvt/configuration_pvt.py
new file mode 100644
index 0000000000000000000000000000000000000000..c97c2703ef2358a38e082aebc01a41e25cc8f5bf
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pvt/configuration_pvt.py
@@ -0,0 +1,162 @@
+# coding=utf-8
+# Copyright 2023 Authors: Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan,
+# Kaitao Song, Ding Liang, Tong Lu, Ping Luo, Ling Shao 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.
+"""Pvt model configuration"""
+
+from collections import OrderedDict
+from typing import Callable, List, Mapping
+
+from packaging import version
+
+from ...configuration_utils import PretrainedConfig
+from ...onnx import OnnxConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class PvtConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`PvtModel`]. It is used to instantiate an Pvt
+ 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 Pvt
+ [Xrenya/pvt-tiny-224](https://huggingface.co/Xrenya/pvt-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:
+ image_size (`int`, *optional*, defaults to 224):
+ The input image size
+ 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.
+ sequence_reduction_ratios (`List[int]`, *optional*, defaults to `[8, 4, 2, 1]`):
+ Sequence reduction ratios in each encoder block.
+ hidden_sizes (`List[int]`, *optional*, defaults to `[64, 128, 320, 512]`):
+ Dimension of each of the encoder blocks.
+ patch_sizes (`List[int]`, *optional*, defaults to `[4, 2, 2, 2]`):
+ 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 `[8, 8, 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.0):
+ 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.
+ qkv_bias (`bool`, *optional*, defaults to `True`):
+ Whether or not a learnable bias should be added to the queries, keys and values.
+ num_labels ('int', *optional*, defaults to 1000):
+ The number of classes.
+ Example:
+
+ ```python
+ >>> from transformers import PvtModel, PvtConfig
+
+ >>> # Initializing a PVT Xrenya/pvt-tiny-224 style configuration
+ >>> configuration = PvtConfig()
+
+ >>> # Initializing a model from the Xrenya/pvt-tiny-224 style configuration
+ >>> model = PvtModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "pvt"
+
+ def __init__(
+ self,
+ image_size: int = 224,
+ num_channels: int = 3,
+ num_encoder_blocks: int = 4,
+ depths: List[int] = [2, 2, 2, 2],
+ sequence_reduction_ratios: List[int] = [8, 4, 2, 1],
+ hidden_sizes: List[int] = [64, 128, 320, 512],
+ patch_sizes: List[int] = [4, 2, 2, 2],
+ strides: List[int] = [4, 2, 2, 2],
+ num_attention_heads: List[int] = [1, 2, 5, 8],
+ mlp_ratios: List[int] = [8, 8, 4, 4],
+ hidden_act: Mapping[str, Callable] = "gelu",
+ hidden_dropout_prob: float = 0.0,
+ attention_probs_dropout_prob: float = 0.0,
+ initializer_range: float = 0.02,
+ drop_path_rate: float = 0.0,
+ layer_norm_eps: float = 1e-6,
+ qkv_bias: bool = True,
+ num_labels: int = 1000,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.image_size = image_size
+ self.num_channels = num_channels
+ self.num_encoder_blocks = num_encoder_blocks
+ self.depths = depths
+ self.sequence_reduction_ratios = sequence_reduction_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.num_labels = num_labels
+ self.qkv_bias = qkv_bias
+
+
+class PvtOnnxConfig(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
+
+ @property
+ def default_onnx_opset(self) -> int:
+ return 12
+
+
+__all__ = ["PvtConfig", "PvtOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/image_processing_pvt.py b/lib/python3.12/site-packages/transformers/models/pvt/image_processing_pvt.py
new file mode 100644
index 0000000000000000000000000000000000000000..75fcba6ea14f0e46146be17deed3b401c14d8440
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pvt/image_processing_pvt.py
@@ -0,0 +1,276 @@
+# 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 Pvt."""
+
+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_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ infer_channel_dimension_format,
+ is_scaled_image,
+ make_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class PvtImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a PVT 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 `Resampling.BILINEAR`):
+ Resampling filter to use if resizing the image. Can be overridden by the `resample` 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.
+ 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`):
+ 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 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_DEFAULT_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.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 {"height": 224, "width": 224}
+ size = get_size_dict(size)
+ self.do_resize = do_resize
+ self.do_rescale = do_rescale
+ self.do_normalize = do_normalize
+ 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
+
+ # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize
+ 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.
+ """
+ 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,
+ )
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = 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,
+ ):
+ """
+ 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_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 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
+ 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)
+
+ 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 do_rescale and is_scaled_image(images[0]):
+ 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_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)
+
+
+__all__ = ["PvtImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/pvt/modeling_pvt.py b/lib/python3.12/site-packages/transformers/models/pvt/modeling_pvt.py
new file mode 100644
index 0000000000000000000000000000000000000000..c061a8fdf0f966d458f0e0dc0c1f0c8eb05fe8ef
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/pvt/modeling_pvt.py
@@ -0,0 +1,671 @@
+# coding=utf-8
+# Copyright 2023 Authors: Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan,
+# Kaitao Song, Ding Liang, Tong Lu, Ping Luo, Ling Shao 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 PVT model."""
+
+import collections
+import math
+from typing import Iterable, 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_outputs import BaseModelOutput, 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,
+)
+from .configuration_pvt import PvtConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "PvtConfig"
+
+_CHECKPOINT_FOR_DOC = "Zetatech/pvt-tiny-224"
+_EXPECTED_OUTPUT_SHAPE = [1, 50, 512]
+
+_IMAGE_CLASS_CHECKPOINT = "Zetatech/pvt-tiny-224"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+# 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.convnext.modeling_convnext.ConvNextDropPath with ConvNext->Pvt
+class PvtDropPath(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 PvtPatchEmbeddings(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: PvtConfig,
+ image_size: Union[int, Iterable[int]],
+ patch_size: Union[int, Iterable[int]],
+ stride: int,
+ num_channels: int,
+ hidden_size: int,
+ cls_token: bool = False,
+ ):
+ super().__init__()
+ self.config = config
+ 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.position_embeddings = nn.Parameter(
+ torch.randn(1, num_patches + 1 if cls_token else num_patches, hidden_size)
+ )
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size)) if cls_token else None
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=stride, stride=patch_size)
+ self.layer_norm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ num_patches = height * width
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == self.config.image_size * self.config.image_size:
+ return self.position_embeddings
+ embeddings = embeddings.reshape(1, height, width, -1).permute(0, 3, 1, 2)
+ interpolated_embeddings = F.interpolate(embeddings, size=(height, width), mode="bilinear")
+ interpolated_embeddings = interpolated_embeddings.reshape(1, -1, height * width).permute(0, 2, 1)
+ return interpolated_embeddings
+
+ def forward(self, pixel_values: torch.Tensor) -> Tuple[torch.Tensor, int, int]:
+ 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."
+ )
+ patch_embed = self.projection(pixel_values)
+ *_, height, width = patch_embed.shape
+ patch_embed = patch_embed.flatten(2).transpose(1, 2)
+ embeddings = self.layer_norm(patch_embed)
+ if self.cls_token is not None:
+ cls_token = self.cls_token.expand(batch_size, -1, -1)
+ embeddings = torch.cat((cls_token, embeddings), dim=1)
+ position_embeddings = self.interpolate_pos_encoding(self.position_embeddings[:, 1:], height, width)
+ position_embeddings = torch.cat((self.position_embeddings[:, :1], position_embeddings), dim=1)
+ else:
+ position_embeddings = self.interpolate_pos_encoding(self.position_embeddings, height, width)
+ embeddings = self.dropout(embeddings + position_embeddings)
+
+ return embeddings, height, width
+
+
+class PvtSelfOutput(nn.Module):
+ def __init__(self, config: PvtConfig, hidden_size: int):
+ super().__init__()
+ self.dense = nn.Linear(hidden_size, 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
+
+
+class PvtEfficientSelfAttention(nn.Module):
+ """Efficient self-attention mechanism with reduction of the sequence [PvT paper](https://arxiv.org/abs/2102.12122)."""
+
+ def __init__(
+ self, config: PvtConfig, hidden_size: int, num_attention_heads: int, sequences_reduction_ratio: float
+ ):
+ 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, bias=config.qkv_bias)
+ self.key = nn.Linear(self.hidden_size, self.all_head_size, bias=config.qkv_bias)
+ self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=config.qkv_bias)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.sequences_reduction_ratio = sequences_reduction_ratio
+ if sequences_reduction_ratio > 1:
+ self.sequence_reduction = nn.Conv2d(
+ hidden_size, hidden_size, kernel_size=sequences_reduction_ratio, stride=sequences_reduction_ratio
+ )
+ self.layer_norm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+
+ def transpose_for_scores(self, hidden_states: int) -> torch.Tensor:
+ 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: torch.Tensor,
+ height: int,
+ width: int,
+ output_attentions: bool = False,
+ ) -> Tuple[torch.Tensor]:
+ query_layer = self.transpose_for_scores(self.query(hidden_states))
+
+ if self.sequences_reduction_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.sequence_reduction(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
+
+
+class PvtAttention(nn.Module):
+ def __init__(
+ self, config: PvtConfig, hidden_size: int, num_attention_heads: int, sequences_reduction_ratio: float
+ ):
+ super().__init__()
+ self.self = PvtEfficientSelfAttention(
+ config,
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ sequences_reduction_ratio=sequences_reduction_ratio,
+ )
+ self.output = PvtSelfOutput(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: torch.Tensor, height: int, width: int, output_attentions: bool = False
+ ) -> Tuple[torch.Tensor]:
+ self_outputs = self.self(hidden_states, height, width, output_attentions)
+
+ attention_output = self.output(self_outputs[0])
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class PvtFFN(nn.Module):
+ def __init__(
+ self,
+ config: PvtConfig,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ ):
+ super().__init__()
+ out_features = out_features if out_features is not None else in_features
+ self.dense1 = nn.Linear(in_features, 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: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense1(hidden_states)
+ 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
+
+
+class PvtLayer(nn.Module):
+ def __init__(
+ self,
+ config: PvtConfig,
+ hidden_size: int,
+ num_attention_heads: int,
+ drop_path: float,
+ sequences_reduction_ratio: float,
+ mlp_ratio: float,
+ ):
+ super().__init__()
+ self.layer_norm_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+ self.attention = PvtAttention(
+ config=config,
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ sequences_reduction_ratio=sequences_reduction_ratio,
+ )
+ self.drop_path = PvtDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+ self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
+ mlp_hidden_size = int(hidden_size * mlp_ratio)
+ self.mlp = PvtFFN(config=config, in_features=hidden_size, hidden_features=mlp_hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor, height: int, width: int, output_attentions: bool = False):
+ self_attention_outputs = self.attention(
+ hidden_states=self.layer_norm_1(hidden_states),
+ height=height,
+ width=width,
+ output_attentions=output_attentions,
+ )
+ attention_output = self_attention_outputs[0]
+ outputs = self_attention_outputs[1:]
+
+ attention_output = self.drop_path(attention_output)
+ hidden_states = attention_output + hidden_states
+
+ mlp_output = self.mlp(self.layer_norm_2(hidden_states))
+
+ mlp_output = self.drop_path(mlp_output)
+ layer_output = hidden_states + mlp_output
+
+ outputs = (layer_output,) + outputs
+
+ return outputs
+
+
+class PvtEncoder(nn.Module):
+ def __init__(self, config: PvtConfig):
+ super().__init__()
+ self.config = config
+
+ # stochastic depth decay rule
+ drop_path_decays = torch.linspace(0, config.drop_path_rate, sum(config.depths)).tolist()
+
+ # patch embeddings
+ embeddings = []
+
+ for i in range(config.num_encoder_blocks):
+ embeddings.append(
+ PvtPatchEmbeddings(
+ config=config,
+ image_size=config.image_size if i == 0 else self.config.image_size // (2 ** (i + 1)),
+ 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],
+ cls_token=i == config.num_encoder_blocks - 1,
+ )
+ )
+ 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(
+ PvtLayer(
+ config=config,
+ hidden_size=config.hidden_sizes[i],
+ num_attention_heads=config.num_attention_heads[i],
+ drop_path=drop_path_decays[cur + j],
+ sequences_reduction_ratio=config.sequence_reduction_ratios[i],
+ mlp_ratio=config.mlp_ratios[i],
+ )
+ )
+ blocks.append(nn.ModuleList(layers))
+
+ self.block = nn.ModuleList(blocks)
+
+ # Layer norms
+ self.layer_norm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ output_attentions: Optional[bool] = False,
+ output_hidden_states: Optional[bool] = False,
+ return_dict: Optional[bool] = True,
+ ) -> Union[Tuple, BaseModelOutput]:
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ batch_size = pixel_values.shape[0]
+ num_blocks = len(self.block)
+ hidden_states = pixel_values
+ for idx, (embedding_layer, block_layer) in enumerate(zip(self.patch_embeddings, self.block)):
+ # first, obtain patch embeddings
+ hidden_states, height, width = embedding_layer(hidden_states)
+ # second, send embeddings through blocks
+ for block in block_layer:
+ layer_outputs = block(hidden_states, height, width, 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 idx != num_blocks - 1:
+ hidden_states = hidden_states.reshape(batch_size, height, width, -1).permute(0, 3, 1, 2).contiguous()
+ hidden_states = self.layer_norm(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, 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 PvtPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = PvtConfig
+ base_model_prefix = "pvt"
+ main_input_name = "pixel_values"
+ _no_split_modules = []
+
+ def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
+ """Initialize the weights"""
+ if isinstance(module, nn.Linear):
+ # 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, 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, PvtPatchEmbeddings):
+ module.position_embeddings.data = nn.init.trunc_normal_(
+ module.position_embeddings.data,
+ mean=0.0,
+ std=self.config.initializer_range,
+ )
+ if module.cls_token is not None:
+ module.cls_token.data = nn.init.trunc_normal_(
+ module.cls_token.data,
+ mean=0.0,
+ std=self.config.initializer_range,
+ )
+
+
+PVT_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 ([`~PvtConfig`]): 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.
+"""
+
+PVT_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 [`PvtImageProcessor.__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 Pvt encoder outputting raw hidden-states without any specific head on top.",
+ PVT_START_DOCSTRING,
+)
+class PvtModel(PvtPreTrainedModel):
+ def __init__(self, config: PvtConfig):
+ super().__init__(config)
+ self.config = config
+
+ # hierarchical Transformer encoder
+ self.encoder = PvtEncoder(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(PVT_INPUTS_DOCSTRING.format("(batch_size, channels, height, width)"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=BaseModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ modality="vision",
+ expected_output=_EXPECTED_OUTPUT_SHAPE,
+ )
+ 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=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,
+ )
+
+
+@add_start_docstrings(
+ """
+ Pvt 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.
+ """,
+ PVT_START_DOCSTRING,
+)
+class PvtForImageClassification(PvtPreTrainedModel):
+ def __init__(self, config: PvtConfig) -> None:
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.pvt = PvtModel(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(PVT_INPUTS_DOCSTRING.format("(batch_size, channels, height, width)"))
+ @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],
+ 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.pvt(
+ pixel_values=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[:, 0, :])
+
+ 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,
+ )
+
+
+__all__ = ["PvtForImageClassification", "PvtModel", "PvtPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__init__.py b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a9f44a7a05fc65136552948be9acef53e5ea106
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 The Qwen 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_qwen2_5_vl import *
+ from .modeling_qwen2_5_vl import *
+ from .processing_qwen2_5_vl import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8320e71c4b468365e96c061be9cd1ba3b16b4b46
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/configuration_qwen2_5_vl.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/configuration_qwen2_5_vl.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e3634b209bb78deaf68ee2512d1ae54298a3b209
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/configuration_qwen2_5_vl.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/modular_qwen2_5_vl.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/modular_qwen2_5_vl.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f91138bef6032015d28f355bc535252bd649c206
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/modular_qwen2_5_vl.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/processing_qwen2_5_vl.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/processing_qwen2_5_vl.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1355fdf7049d86dbfe2c8947a914862ed0f6f7b5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/__pycache__/processing_qwen2_5_vl.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed3505728a6331220a60ca71aa6e1dba7aef09a0
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py
@@ -0,0 +1,258 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen2_5_vl.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen Team 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.
+from ...configuration_utils import PretrainedConfig
+from ...modeling_rope_utils import rope_config_validation
+
+
+class Qwen2_5_VLVisionConfig(PretrainedConfig):
+ model_type = "qwen2_5_vl"
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ depth=32,
+ hidden_size=3584,
+ hidden_act="silu",
+ intermediate_size=3420,
+ num_heads=16,
+ in_channels=3,
+ patch_size=14,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+ tokens_per_second=4,
+ window_size=112,
+ out_hidden_size=3584,
+ fullatt_block_indexes=[7, 15, 23, 31],
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.depth = depth
+ self.hidden_size = hidden_size
+ self.hidden_act = hidden_act
+ self.intermediate_size = intermediate_size
+ self.num_heads = num_heads
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.spatial_merge_size = spatial_merge_size
+ self.temporal_patch_size = temporal_patch_size
+ self.tokens_per_second = tokens_per_second
+ self.window_size = window_size
+ self.fullatt_block_indexes = fullatt_block_indexes
+ self.out_hidden_size = out_hidden_size
+
+
+class Qwen2_5_VLConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
+ Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
+ with the defaults will yield a similar configuration to that of
+ Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
+
+ 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 152064):
+ Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`Qwen2_5_VLModel`]
+ hidden_size (`int`, *optional*, defaults to 8192):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 29568):
+ Dimension of the MLP representations.
+ num_hidden_layers (`int`, *optional*, defaults to 80):
+ 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.
+ num_key_value_heads (`int`, *optional*, defaults to 8):
+ 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 `32`.
+ 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 32768):
+ 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-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`.
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether the model's input and output word embeddings should be tied.
+ rope_theta (`float`, *optional*, defaults to 1000000.0):
+ The base period of the RoPE embeddings.
+ use_sliding_window (`bool`, *optional*, defaults to `False`):
+ Whether to use sliding window attention.
+ sliding_window (`int`, *optional*, defaults to 4096):
+ Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+ max_window_layers (`int`, *optional*, defaults to 80):
+ The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+ vision_config (`Dict`, *optional*):
+ The config for the visual encoder initialization.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+ accordingly.
+ Expected contents:
+ `rope_type` (`str`):
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+ 'llama3'], with 'default' being the original RoPE implementation.
+ `factor` (`float`, *optional*):
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+ original maximum pre-trained length.
+ `original_max_position_embeddings` (`int`, *optional*):
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+ pretraining.
+ `attention_factor` (`float`, *optional*):
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
+ `factor` field to infer the suggested value.
+ `beta_fast` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+ ramp function. If unspecified, it defaults to 32.
+ `beta_slow` (`float`, *optional*):
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+ ramp function. If unspecified, it defaults to 1.
+ `short_factor` (`List[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `long_factor` (`List[float]`, *optional*):
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+ size divided by the number of attention heads divided by 2
+ `low_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+ `high_freq_factor` (`float`, *optional*):
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+
+ ```python
+ >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
+
+ >>> # Initializing a Qwen2_5_VL style configuration
+ >>> configuration = Qwen2_5_VLConfig()
+
+ >>> # Initializing a model from the Qwen2-VL-7B style configuration
+ >>> model = Qwen2_5_VLForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_vl"
+ sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
+ keys_to_ignore_at_inference = ["past_key_values"]
+ # Default tensor parallel plan for base model `Qwen2_5_VL`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ def __init__(
+ self,
+ vocab_size=152064,
+ hidden_size=8192,
+ intermediate_size=29568,
+ num_hidden_layers=80,
+ num_attention_heads=64,
+ num_key_value_heads=8,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-05,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=1000000.0,
+ use_sliding_window=False,
+ sliding_window=4096,
+ max_window_layers=80,
+ attention_dropout=0.0,
+ vision_config=None,
+ rope_scaling=None,
+ **kwargs,
+ ):
+ if isinstance(vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**vision_config)
+ elif vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"]()
+
+ 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.use_sliding_window = use_sliding_window
+ self.sliding_window = sliding_window
+ self.max_window_layers = max_window_layers
+
+ # 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.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.attention_dropout = attention_dropout
+ self.rope_scaling = rope_scaling
+
+ # Validate the correctness of rotary position embeddings parameters
+ # BC: if there is a 'type' field, move it to 'rope_type'.
+ # and change type from 'mrope' to 'default' because `mrope` does default RoPE calculations
+ # one can set it to "linear"/"dynamic" etc. to have scaled RoPE
+ # TODO: @raushan update config in the hub
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
+ if self.rope_scaling["type"] == "mrope":
+ self.rope_scaling["type"] = "default"
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+ rope_config_validation(self, ignore_keys={"mrope_section"})
+
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["Qwen2_5_VLConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb455444e5cbdaedb013b4bbd3de64de35093bc1
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
@@ -0,0 +1,2035 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen2_5_vl.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen Team 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.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import PreTrainedModel
+from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import apply_rotary_emb, flash_attn_varlen_func
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "Qwen2_5_VLConfig"
+
+
+class Qwen2_5_VLMLP(nn.Module):
+ def __init__(self, config, bias: bool = False):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class Qwen2_5_VisionPatchEmbed(nn.Module):
+ def __init__(
+ self,
+ patch_size: int = 14,
+ temporal_patch_size: int = 2,
+ in_channels: int = 3,
+ embed_dim: int = 1152,
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.temporal_patch_size = temporal_patch_size
+ self.in_channels = in_channels
+ self.embed_dim = embed_dim
+
+ kernel_size = [temporal_patch_size, patch_size, patch_size]
+ self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen2_5_VisionRotaryEmbedding(nn.Module):
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ freqs = torch.outer(seq, self.inv_freq)
+ return freqs
+
+
+class Qwen2RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ Qwen2RMSNorm 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Qwen2_5_VLPatchMerger(nn.Module):
+ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
+ super().__init__()
+ self.hidden_size = context_dim * (spatial_merge_size**2)
+ self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6)
+ self.mlp = nn.Sequential(
+ nn.Linear(self.hidden_size, self.hidden_size),
+ nn.GELU(),
+ nn.Linear(self.hidden_size, dim),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
+ return x
+
+
+def apply_rotary_pos_emb_flashatt(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ cos = cos.chunk(2, dim=-1)[0].contiguous()
+ sin = sin.chunk(2, dim=-1)[0].contiguous()
+ q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
+ k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k)
+ return q_embed, k_embed
+
+
+class Qwen2_5_VLVisionFlashAttention2(nn.Module):
+ def __init__(self, dim: int, num_heads: int = 16) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.proj = nn.Linear(dim, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ if position_embeddings is None:
+ logger.warning_once(
+ "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+ "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed "
+ "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be "
+ "removed and `position_embeddings` will be mandatory."
+ )
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+ else:
+ cos, sin = position_embeddings
+ q, k = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), k.unsqueeze(0), cos, sin)
+ q = q.squeeze(0)
+ k = k.squeeze(0)
+
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+ attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape(
+ seq_length, -1
+ )
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+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_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+class Qwen2_5_VLVisionAttention(nn.Module):
+ def __init__(self, dim: int, num_heads: int = 16) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.proj = nn.Linear(dim, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ if position_embeddings is None:
+ logger.warning_once(
+ "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+ "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed "
+ "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be "
+ "removed and `position_embeddings` will be mandatory."
+ )
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+ else:
+ cos, sin = position_embeddings
+ q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
+
+ attention_mask = torch.full(
+ [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype
+ )
+ for i in range(1, len(cu_seqlens)):
+ attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
+
+ q = q.transpose(0, 1)
+ k = k.transpose(0, 1)
+ v = v.transpose(0, 1)
+ attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
+ attn_weights = attn_weights + attention_mask
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+ attn_output = torch.matmul(attn_weights, v)
+ attn_output = attn_output.transpose(0, 1)
+ attn_output = attn_output.reshape(seq_length, -1)
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen2_5_VLVisionSdpaAttention(nn.Module):
+ def __init__(self, dim: int, num_heads: int = 16) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.proj = nn.Linear(dim, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ if position_embeddings is None:
+ logger.warning_once(
+ "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+ "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed "
+ "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be "
+ "removed and `position_embeddings` will be mandatory."
+ )
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+ else:
+ cos, sin = position_embeddings
+ q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
+
+ attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool)
+ for i in range(1, len(cu_seqlens)):
+ attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = True
+ q = q.transpose(0, 1)
+ k = k.transpose(0, 1)
+ v = v.transpose(0, 1)
+ attn_output = F.scaled_dot_product_attention(
+ q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0
+ )
+ attn_output = attn_output.squeeze(0).transpose(0, 1)
+ attn_output = attn_output.reshape(seq_length, -1)
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+QWEN2_5_VL_VISION_ATTENTION_CLASSES = {
+ "eager": Qwen2_5_VLVisionAttention,
+ "flash_attention_2": Qwen2_5_VLVisionFlashAttention2,
+ "sdpa": Qwen2_5_VLVisionSdpaAttention,
+}
+
+
+class Qwen2_5_VLVisionBlock(nn.Module):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+ self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation](
+ config.hidden_size, num_heads=config.num_heads
+ )
+ self.mlp = Qwen2_5_VLMLP(config, bias=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+Qwen2_5_VL_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 ([`Qwen2_5_VLConfig`]):
+ 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 Qwen2_5_VL Model outputting raw hidden-states without any specific head on top.",
+ Qwen2_5_VL_START_DOCSTRING,
+)
+class Qwen2_5_VLPreTrainedModel(PreTrainedModel):
+ config_class = Qwen2_5_VLConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_cache_class = True
+ _supports_static_cache = False # TODO (joao): fix. torch.compile failing probably due to `cache_positions`
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv3d)):
+ 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_()
+
+
+class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):
+ config_class = Qwen2_5_VLVisionConfig
+ _no_split_modules = ["Qwen2_5_VLVisionBlock"]
+
+ def __init__(self, config, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.fullatt_block_indexes = config.fullatt_block_indexes
+ self.window_size = config.window_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen2_5_VisionPatchEmbed(
+ patch_size=config.patch_size,
+ temporal_patch_size=config.temporal_patch_size,
+ in_channels=config.in_channels,
+ embed_dim=config.hidden_size,
+ )
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList(
+ [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)]
+ )
+ self.merger = Qwen2_5_VLPatchMerger(
+ dim=config.out_hidden_size,
+ context_dim=config.hidden_size,
+ spatial_merge_size=config.spatial_merge_size,
+ )
+ self.gradient_checkpointing = False
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ max_grid_size = grid_thw[:, 1:].max()
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+ return rotary_pos_emb
+
+ def get_window_index(self, grid_thw):
+ window_index: list = []
+ cu_window_seqlens: list = [0]
+ window_index_id = 0
+ vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size
+
+ for grid_t, grid_h, grid_w in grid_thw:
+ llm_grid_h, llm_grid_w = (
+ grid_h // self.spatial_merge_size,
+ grid_w // self.spatial_merge_size,
+ )
+ index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
+ pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
+ pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
+ num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
+ num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+ index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)
+ index_padded = index_padded.reshape(
+ grid_t,
+ num_windows_h,
+ vit_merger_window_size,
+ num_windows_w,
+ vit_merger_window_size,
+ )
+ index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(
+ grid_t,
+ num_windows_h * num_windows_w,
+ vit_merger_window_size,
+ vit_merger_window_size,
+ )
+ seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)
+ index_padded = index_padded.reshape(-1)
+ index_new = index_padded[index_padded != -100]
+ window_index.append(index_new + window_index_id)
+ cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]
+ cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
+ window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+ window_index = torch.cat(window_index, dim=0)
+
+ return window_index, cu_window_seqlens
+
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+ window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+ cu_window_seqlens = torch.tensor(
+ cu_window_seqlens,
+ device=hidden_states.device,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ hidden_states = hidden_states[window_index, :, :]
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ for layer_num, blk in enumerate(self.blocks):
+ if layer_num in self.fullatt_block_indexes:
+ cu_seqlens_now = cu_seqlens
+ else:
+ cu_seqlens_now = cu_window_seqlens
+ if self.gradient_checkpointing and self.training:
+ hidden_states = self._gradient_checkpointing_func(
+ blk.__call__, hidden_states, cu_seqlens_now, None, position_embeddings
+ )
+ else:
+ hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens_now, position_embeddings=position_embeddings)
+
+ hidden_states = self.merger(hidden_states)
+ reverse_indices = torch.argsort(window_index)
+ hidden_states = hidden_states[reverse_indices, :]
+
+ return hidden_states
+
+
+class Qwen2_5_VLRotaryEmbedding(nn.Module):
+ def __init__(self, config: Qwen2_5_VLConfig, device=None):
+ super().__init__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, Qwen2_5_VL has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Qwen2MLP(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]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
+
+ Explanation:
+ Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding
+ sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For
+ vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.
+ Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.
+ For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,
+ height and width) of text embedding is always the same, so the text embedding rotary position embedding has no
+ difference with modern LLMs.
+
+ 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.
+ mrope_section(`List(int)`):
+ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.
+ 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.
+ """
+ mrope_section = mrope_section * 2
+ cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+ unsqueeze_dim
+ )
+ sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+ unsqueeze_dim
+ )
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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 Qwen2_5_VLAttention(nn.Module):
+ """
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
+ and "Generating Long Sequences with Sparse Transformers".
+ """
+
+ def __init__(self, config: Qwen2_5_VLConfig, 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 `layer_idx` is not recommended and will "
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ 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.is_causal = True
+ self.attention_dropout = config.attention_dropout
+ self.rope_scaling = config.rope_scaling
+
+ 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=True)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+ self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config)
+
+ 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,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ ) -> 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, -1, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_multimodal_rotary_pos_emb(
+ query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
+ )
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # repeat k/v heads if n_kv_heads < n_heads
+ 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
+
+ # Fix precision issues in Qwen2-VL float16 inference
+ # Replace inf values with zeros in attention weights to prevent NaN propagation
+ if query_states.dtype == torch.float16:
+ attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)
+
+ # 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, -1)
+
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention):
+ """
+ Qwen2_5_VL flash attention module, following Qwen2_5_VL attention module. This module inherits from `Qwen2_5_VLAttention`
+ 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. Additionally, for sliding window attention, we apply SWA only to the bottom
+ config.max_window_layers layers.
+ """
+
+ 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 alignment, 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 = flash_attn_supports_top_left_mask()
+
+ 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,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ ):
+ 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, -1, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
+ cos, sin = position_embeddings
+ query_states, key_states = apply_multimodal_rotary_pos_emb(
+ query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
+ )
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # repeat k/v heads if n_kv_heads < n_heads
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+ # 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_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)
+
+ # Reashape to the expected shape for Flash Attention
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+
+ if (
+ self.config.use_sliding_window
+ and getattr(self.config, "sliding_window", None) is not None
+ and self.layer_idx >= self.config.max_window_layers
+ ):
+ sliding_window = self.config.sliding_window
+ else:
+ sliding_window = None
+
+ attn_output = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=dropout_rate,
+ sliding_window=sliding_window,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ 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
+
+
+class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention):
+ """
+ Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+ `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+ SDPA API.
+ """
+
+ # Adapted from Qwen2Attention.forward
+ 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,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ ) -> 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(
+ "Qwen2_5_VLModel is using Qwen2_5_VLSdpaAttention, 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,
+ position_embeddings=position_embeddings,
+ )
+
+ 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, -1, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_multimodal_rotary_pos_emb(
+ query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
+ )
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
+ 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: # no matter the length, we just slice it
+ causal_mask = attention_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 attention_mask is not None:
+ query_states = query_states.contiguous()
+ key_states = key_states.contiguous()
+ value_states = value_states.contiguous()
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+ is_causal = True if causal_mask is None and q_len > 1 else False
+
+ 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=is_causal,
+ )
+
+ 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
+
+
+QWEN2_5_VL_ATTENTION_CLASSES = {
+ "eager": Qwen2_5_VLAttention,
+ "flash_attention_2": Qwen2_5_VLFlashAttention2,
+ "sdpa": Qwen2_5_VLSdpaAttention,
+}
+
+
+class Qwen2_5_VLDecoderLayer(nn.Module):
+ def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ if config.use_sliding_window and config._attn_implementation != "flash_attention_2":
+ logger.warning_once(
+ f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
+ "unexpected results may be encountered."
+ )
+ self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+
+ self.mlp = Qwen2MLP(config)
+ self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen2RMSNorm(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,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
+ **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, sequence_length)` where padding elements are indicated by 0.
+ 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
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+
+ 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,
+ position_embeddings=position_embeddings,
+ )
+ 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
+
+
+@add_start_docstrings(
+ "The bare Qwen2_5_VL Model outputting raw hidden-states without any specific head on top.",
+ Qwen2_5_VL_START_DOCSTRING,
+)
+class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel):
+ def __init__(self, config: Qwen2_5_VLConfig):
+ 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(
+ [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self._attn_implementation = config._attn_implementation
+ self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config)
+
+ 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
+
+ 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[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 must specify exactly one of input_ids or inputs_embeds")
+
+ 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
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache()
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = torch.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+ )
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.dim() == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ causal_mask = self._update_causal_mask(
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+ )
+
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # 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,
+ position_embeddings,
+ )
+ 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,
+ position_embeddings=position_embeddings,
+ )
+
+ 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,
+ )
+
+ def _update_causal_mask(
+ self,
+ attention_mask: torch.Tensor,
+ input_tensor: torch.Tensor,
+ cache_position: torch.Tensor,
+ past_key_values: Cache,
+ output_attentions: bool = False,
+ ):
+ if self.config._attn_implementation == "flash_attention_2":
+ if attention_mask is not None and past_key_values is not None:
+ is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
+ if is_padding_right:
+ raise ValueError(
+ "You are attempting to perform batched generation with padding_side='right'"
+ " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to "
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
+ )
+ if attention_mask is not None and 0.0 in attention_mask:
+ return attention_mask
+ return None
+
+ # 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. This feature is not compatible with static cache, as SDPA will fail
+ # to infer the attention mask.
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ using_static_cache = isinstance(past_key_values, StaticCache)
+ using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+ if (
+ self.config._attn_implementation == "sdpa"
+ and not (using_static_cache or using_sliding_window_cache)
+ and not output_attentions
+ ):
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
+ attention_mask,
+ inputs_embeds=input_tensor,
+ past_key_values_length=past_seen_tokens,
+ sliding_window=self.config.sliding_window,
+ is_training=self.training,
+ ):
+ return None
+
+ dtype, device = input_tensor.dtype, input_tensor.device
+ min_dtype = torch.finfo(dtype).min
+ sequence_length = input_tensor.shape[1]
+ # SlidingWindowCache or StaticCache
+ if using_sliding_window_cache or using_static_cache:
+ target_length = past_key_values.get_max_cache_shape()
+ # DynamicCache or no cache
+ else:
+ target_length = (
+ attention_mask.shape[-1]
+ if isinstance(attention_mask, torch.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask,
+ sequence_length=sequence_length,
+ target_length=target_length,
+ dtype=dtype,
+ device=device,
+ cache_position=cache_position,
+ batch_size=input_tensor.shape[0],
+ config=self.config,
+ past_key_values=past_key_values,
+ )
+
+ if (
+ self.config._attn_implementation == "sdpa"
+ and attention_mask is not None
+ and attention_mask.device.type in ["cuda", "xpu"]
+ and not output_attentions
+ ):
+ # 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
+
+ @staticmethod
+ def _prepare_4d_causal_attention_mask_with_cache_position(
+ attention_mask: torch.Tensor,
+ sequence_length: int,
+ target_length: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ cache_position: torch.Tensor,
+ batch_size: int,
+ config: Qwen2_5_VLConfig,
+ past_key_values: Cache,
+ ):
+ """
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+ sequence_length (`int`):
+ The sequence length being processed.
+ target_length (`int`):
+ The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+ dtype (`torch.dtype`):
+ The dtype to use for the 4D attention mask.
+ device (`torch.device`):
+ The device to place the 4D attention mask on.
+ cache_position (`torch.Tensor`):
+ Indices depicting the position of the input sequence tokens in the sequence.
+ batch_size (`torch.Tensor`):
+ Batch size.
+ config (`Qwen2_5_VLConfig`):
+ The model's configuration class
+ past_key_values (`Cache`):
+ The cache class that is being used currently to generate
+ """
+ if attention_mask is not None and attention_mask.dim() == 4:
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+ causal_mask = attention_mask
+ else:
+ min_dtype = torch.finfo(dtype).min
+ causal_mask = torch.full(
+ (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+ )
+ diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+ if config.sliding_window is not None:
+ # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
+ # the check is needed to verify is current checkpoint was trained with sliding window or not
+ if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
+ sliding_attend_mask = torch.arange(target_length, device=device) <= (
+ cache_position.reshape(-1, 1) - config.sliding_window
+ )
+ diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+ causal_mask *= diagonal_attend_mask
+ causal_mask = causal_mask[None, None, :, :].expand(batch_size, 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.shape[-1] > target_length:
+ attention_mask = attention_mask[:, :target_length]
+ mask_length = attention_mask.shape[-1]
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+ causal_mask.device
+ )
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+ return causal_mask
+
+
+@dataclass
+class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput):
+ """
+ Base class for Qwen2_5_VL causal language model (or autoregressive) outputs.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token 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).
+ 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)`)
+
+ 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.
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+
+
+QWEN2_5_VL_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.
+ pixel_values (`torch.FloatTensor` of shape `(seq_length, num_channels * image_size * image_size)):
+ The tensors corresponding to the input images. Pixel values can be obtained using
+ [`AutoImageProcessor`]. See [`Qwen2_5_VLImageProcessor.__call__`] for details. [`Qwen2_5_VLProcessor`] uses
+ [`Qwen2_5_VLImageProcessor`] for processing images.
+ pixel_values_videos (`torch.FloatTensor` of shape `(seq_length, num_channels * temporal_size * image_size * image_size)):
+ The tensors corresponding to the input videos. Pixel values can be obtained using
+ [`AutoImageProcessor`]. See [`Qwen2_5_VLImageProcessor.__call__`] for details. [`Qwen2_5_VLProcessor`] uses
+ [`Qwen2_5_VLImageProcessor`] for processing videos.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+"""
+
+
+class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+ config_class = Qwen2_5_VLConfig
+ _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
+ self.model = Qwen2_5_VLModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+ self.rope_deltas = None # cache rope_deltas here
+
+ # 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
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ second_per_grid_ts: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ 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.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position 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**.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids in enumerate(total_input_ids):
+ input_ids = input_ids[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = (
+ image_grid_thw[image_index][0],
+ image_grid_thw[image_index][1],
+ image_grid_thw[image_index][2],
+ )
+ second_per_grid_t = 0
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+
+ else:
+ t, h, w = (
+ video_grid_thw[video_index][0],
+ video_grid_thw[video_index][1],
+ video_grid_thw[video_index][2],
+ )
+ if second_per_grid_ts is not None:
+ second_per_grid_t = second_per_grid_ts[video_index]
+ else:
+ second_per_grid_t = 1.0
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ range_tensor = torch.arange(llm_grid_t).view(-1, 1)
+ expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
+
+ time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second
+
+ time_tensor_long = time_tensor.long()
+ t_index = time_tensor_long.flatten()
+
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+ else:
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = (
+ torch.arange(input_ids.shape[1], device=input_ids.device)
+ .view(1, 1, -1)
+ .expand(3, input_ids.shape[0], -1)
+ )
+ mrope_position_deltas = torch.zeros(
+ [input_ids.shape[0], 1],
+ device=input_ids.device,
+ dtype=input_ids.dtype,
+ )
+
+ return position_ids, mrope_position_deltas
+
+ @add_start_docstrings_to_model_forward(QWEN2_5_VL_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Qwen2_5_VLCausalLMOutputWithPast, 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[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,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ rope_deltas: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ second_per_grid_ts: Optional[torch.Tensor] = None,
+ ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
+ 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:
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+ >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image"},
+ {"type": "text", "text": "What is shown in this image?"},
+ ],
+ },
+ ]
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
+
+ >>> # 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]
+ "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+ ```"""
+
+ 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 None:
+ inputs_embeds = self.model.embed_tokens(input_ids)
+ if pixel_values is not None:
+ pixel_values = pixel_values.type(self.visual.dtype)
+ image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+ n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+ n_image_features = image_embeds.shape[0]
+ if n_image_tokens != n_image_features:
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+
+ mask = input_ids == self.config.image_token_id
+ mask_unsqueezed = mask.unsqueeze(-1)
+ mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+ image_mask = mask_expanded.to(inputs_embeds.device)
+
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+ n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
+ n_video_features = video_embeds.shape[0]
+ if n_video_tokens != n_video_features:
+ raise ValueError(
+ f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
+ )
+
+ mask = input_ids == self.config.video_token_id
+ mask_unsqueezed = mask.unsqueeze(-1)
+ mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+ video_mask = mask_expanded.to(inputs_embeds.device)
+
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(inputs_embeds.device)
+
+ # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
+ if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
+ # calculate RoPE index once per generation in the pre-fill stage only
+ if (
+ (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ or (past_key_values is None or past_key_values.get_seq_length() == 0)
+ ):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ second_per_grid_ts,
+ attention_mask,
+ )
+ self.rope_deltas = rope_deltas
+ # then use the prev pre-calculated rope-deltas to get the correct position ids
+ else:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ delta = (
+ (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+ if cache_position is not None
+ else 0
+ )
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ if cache_position is not None: # otherwise `deltas` is an int `0`
+ delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ 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)
+
+ loss = None
+ if labels is not None:
+ # Upcast to float if we need to compute the loss to avoid potential precision issues
+ logits = logits.float()
+ # 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 Qwen2_5_VLCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ cache_position=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ second_per_grid_ts=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ second_per_grid_ts=second_per_grid_ts,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
+ model_inputs["position_ids"] = None
+
+ if cache_position[0] != 0:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+ def _get_image_nums_and_video_nums(
+ self,
+ input_ids: Optional[torch.LongTensor],
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+ These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Returns:
+ image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+ video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+ """
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+
+ vision_start_mask = input_ids == vision_start_token_id
+ vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
+ image_mask = input_ids == image_token_id
+ video_mask = input_ids == video_token_id
+ image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
+ video_nums = torch.sum(vision_first_mask & video_mask, dim=1)
+
+ return image_nums, video_nums
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: Optional[torch.LongTensor] = None,
+ **model_kwargs,
+ ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
+ # Overwritten -- Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ video_grid_thw = model_kwargs.get("video_grid_thw", None)
+ image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids)
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # split images into samples
+ samples = torch.split(image_grid_thw, list(image_nums))
+ # compute the sequence length of images for each sample
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # get the num of images for each sample
+ lengths = list(image_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "pixel_values_videos":
+ samples = torch.split(video_grid_thw, list(video_nums))
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "video_grid_thw":
+ lengths = list(video_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "second_per_grid_ts":
+ if not isinstance(dict_to_expand[key], list):
+ raise TypeError(
+ f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead."
+ )
+ tensor = torch.tensor(dict_to_expand[key])
+ lengths = list(video_nums)
+ tensor = _repeat_interleave_samples(tensor, lengths=lengths, repeat_times=expand_size)
+ dict_to_expand[key] = tensor.tolist()
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if (
+ key != "cache_position"
+ and dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ # input_ids is required for expanding visual inputs
+ # If input_ids is unavailable, visual inputs will not be used; therefore, there is no need to expand visual inputs.
+ if input_ids is not None and input_ids.numel() != 0:
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
+
+
+__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..925d59df2fa1f542b2940ecf5714dddb0a294b19
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py
@@ -0,0 +1,968 @@
+# coding=utf-8
+# Copyright 2025 The Qwen Team 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 Qwen2.5-VL model."""
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch.nn import CrossEntropyLoss
+
+from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig
+from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+ PatchEmbed,
+ PatchMerger,
+ Qwen2RMSNorm,
+ Qwen2VLCausalLMOutputWithPast,
+ Qwen2VLForConditionalGeneration,
+ Qwen2VLModel,
+ Qwen2VLPreTrainedModel,
+ VisionAttention,
+ VisionRotaryEmbedding,
+ VisionSdpaAttention,
+)
+from transformers.models.qwen2_vl.processing_qwen2_vl import Qwen2VLImagesKwargs, Qwen2VLProcessor
+
+from ...activations import ACT2FN
+from ...configuration_utils import PretrainedConfig
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, VideoInput
+from ...modeling_flash_attention_utils import is_flash_attn_available
+from ...processing_utils import ProcessingKwargs, Unpack, VideosKwargs
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import logging
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import apply_rotary_emb, flash_attn_varlen_func
+
+
+logger = logging.get_logger(__name__)
+
+
+def apply_rotary_pos_emb_flashatt(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ cos = cos.chunk(2, dim=-1)[0].contiguous()
+ sin = sin.chunk(2, dim=-1)[0].contiguous()
+ q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
+ k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k)
+ return q_embed, k_embed
+
+
+class Qwen2_5_VLVisionConfig(PretrainedConfig):
+ model_type = "qwen2_5_vl"
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ depth=32,
+ hidden_size=3584,
+ hidden_act="silu",
+ intermediate_size=3420,
+ num_heads=16,
+ in_channels=3,
+ patch_size=14,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+ tokens_per_second=4,
+ window_size=112,
+ out_hidden_size=3584,
+ fullatt_block_indexes=[7, 15, 23, 31],
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.depth = depth
+ self.hidden_size = hidden_size
+ self.hidden_act = hidden_act
+ self.intermediate_size = intermediate_size
+ self.num_heads = num_heads
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.spatial_merge_size = spatial_merge_size
+ self.temporal_patch_size = temporal_patch_size
+ self.tokens_per_second = tokens_per_second
+ self.window_size = window_size
+ self.fullatt_block_indexes = fullatt_block_indexes
+ self.out_hidden_size = out_hidden_size
+
+
+class Qwen2_5_VLConfig(Qwen2VLConfig):
+ model_type = "qwen2_5_vl"
+ sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
+
+
+class Qwen2_5_VLMLP(nn.Module):
+ def __init__(self, config, bias: bool = False):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class Qwen2_5_VisionPatchEmbed(PatchEmbed):
+ pass
+
+
+class Qwen2_5_VisionRotaryEmbedding(VisionRotaryEmbedding):
+ pass
+
+
+class Qwen2_5_VLPatchMerger(PatchMerger):
+ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
+ super().__init__(dim, context_dim, spatial_merge_size)
+ self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6)
+
+
+class Qwen2_5_VLVisionFlashAttention2(nn.Module):
+ def __init__(self, dim: int, num_heads: int = 16) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.proj = nn.Linear(dim, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ if position_embeddings is None:
+ logger.warning_once(
+ "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+ "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed "
+ "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be "
+ "removed and `position_embeddings` will be mandatory."
+ )
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+ else:
+ cos, sin = position_embeddings
+ q, k = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), k.unsqueeze(0), cos, sin)
+ q = q.squeeze(0)
+ k = k.squeeze(0)
+
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+ attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape(
+ seq_length, -1
+ )
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen2_5_VLVisionAttention(VisionAttention):
+ pass
+
+
+class Qwen2_5_VLVisionSdpaAttention(VisionSdpaAttention):
+ pass
+
+
+QWEN2_5_VL_VISION_ATTENTION_CLASSES = {
+ "eager": Qwen2_5_VLVisionAttention,
+ "flash_attention_2": Qwen2_5_VLVisionFlashAttention2,
+ "sdpa": Qwen2_5_VLVisionSdpaAttention,
+}
+
+
+class Qwen2_5_VLVisionBlock(nn.Module):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+ self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation](
+ config.hidden_size, num_heads=config.num_heads
+ )
+ self.mlp = Qwen2_5_VLMLP(config, bias=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen2_5_VLPreTrainedModel(Qwen2VLPreTrainedModel):
+ pass
+
+
+class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):
+ config_class = Qwen2_5_VLVisionConfig
+ _no_split_modules = ["Qwen2_5_VLVisionBlock"]
+
+ def __init__(self, config, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.fullatt_block_indexes = config.fullatt_block_indexes
+ self.window_size = config.window_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen2_5_VisionPatchEmbed(
+ patch_size=config.patch_size,
+ temporal_patch_size=config.temporal_patch_size,
+ in_channels=config.in_channels,
+ embed_dim=config.hidden_size,
+ )
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList(
+ [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)]
+ )
+ self.merger = Qwen2_5_VLPatchMerger(
+ dim=config.out_hidden_size,
+ context_dim=config.hidden_size,
+ spatial_merge_size=config.spatial_merge_size,
+ )
+ self.gradient_checkpointing = False
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ max_grid_size = grid_thw[:, 1:].max()
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+ return rotary_pos_emb
+
+ def get_window_index(self, grid_thw):
+ window_index: list = []
+ cu_window_seqlens: list = [0]
+ window_index_id = 0
+ vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size
+
+ for grid_t, grid_h, grid_w in grid_thw:
+ llm_grid_h, llm_grid_w = (
+ grid_h // self.spatial_merge_size,
+ grid_w // self.spatial_merge_size,
+ )
+ index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
+ pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
+ pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
+ num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
+ num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+ index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)
+ index_padded = index_padded.reshape(
+ grid_t,
+ num_windows_h,
+ vit_merger_window_size,
+ num_windows_w,
+ vit_merger_window_size,
+ )
+ index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(
+ grid_t,
+ num_windows_h * num_windows_w,
+ vit_merger_window_size,
+ vit_merger_window_size,
+ )
+ seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)
+ index_padded = index_padded.reshape(-1)
+ index_new = index_padded[index_padded != -100]
+ window_index.append(index_new + window_index_id)
+ cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]
+ cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
+ window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+ window_index = torch.cat(window_index, dim=0)
+
+ return window_index, cu_window_seqlens
+
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+ window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+ cu_window_seqlens = torch.tensor(
+ cu_window_seqlens,
+ device=hidden_states.device,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ hidden_states = hidden_states[window_index, :, :]
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ for layer_num, blk in enumerate(self.blocks):
+ if layer_num in self.fullatt_block_indexes:
+ cu_seqlens_now = cu_seqlens
+ else:
+ cu_seqlens_now = cu_window_seqlens
+ if self.gradient_checkpointing and self.training:
+ hidden_states = self._gradient_checkpointing_func(
+ blk.__call__, hidden_states, cu_seqlens_now, None, position_embeddings
+ )
+ else:
+ hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens_now, position_embeddings=position_embeddings)
+
+ hidden_states = self.merger(hidden_states)
+ reverse_indices = torch.argsort(window_index)
+ hidden_states = hidden_states[reverse_indices, :]
+
+ return hidden_states
+
+
+class Qwen2_5_VLModel(Qwen2VLModel):
+ pass
+
+
+@dataclass
+class Qwen2_5_VLCausalLMOutputWithPast(Qwen2VLCausalLMOutputWithPast):
+ pass
+
+
+class Qwen2_5_VLForConditionalGeneration(Qwen2VLForConditionalGeneration):
+ config_class = Qwen2_5_VLConfig
+ _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ second_per_grid_ts: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ 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.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position 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**.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids in enumerate(total_input_ids):
+ input_ids = input_ids[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = (
+ image_grid_thw[image_index][0],
+ image_grid_thw[image_index][1],
+ image_grid_thw[image_index][2],
+ )
+ second_per_grid_t = 0
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+
+ else:
+ t, h, w = (
+ video_grid_thw[video_index][0],
+ video_grid_thw[video_index][1],
+ video_grid_thw[video_index][2],
+ )
+ if second_per_grid_ts is not None:
+ second_per_grid_t = second_per_grid_ts[video_index]
+ else:
+ second_per_grid_t = 1.0
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ range_tensor = torch.arange(llm_grid_t).view(-1, 1)
+ expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
+
+ time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second
+
+ time_tensor_long = time_tensor.long()
+ t_index = time_tensor_long.flatten()
+
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+ else:
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = (
+ torch.arange(input_ids.shape[1], device=input_ids.device)
+ .view(1, 1, -1)
+ .expand(3, input_ids.shape[0], -1)
+ )
+ mrope_position_deltas = torch.zeros(
+ [input_ids.shape[0], 1],
+ device=input_ids.device,
+ dtype=input_ids.dtype,
+ )
+
+ return position_ids, mrope_position_deltas
+
+ 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[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,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ rope_deltas: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ second_per_grid_ts: Optional[torch.Tensor] = None,
+ ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
+ 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:
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+ >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image"},
+ {"type": "text", "text": "What is shown in this image?"},
+ ],
+ },
+ ]
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
+
+ >>> # 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]
+ "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+ ```"""
+
+ 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 None:
+ inputs_embeds = self.model.embed_tokens(input_ids)
+ if pixel_values is not None:
+ pixel_values = pixel_values.type(self.visual.dtype)
+ image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+ n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+ n_image_features = image_embeds.shape[0]
+ if n_image_tokens != n_image_features:
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+
+ mask = input_ids == self.config.image_token_id
+ mask_unsqueezed = mask.unsqueeze(-1)
+ mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+ image_mask = mask_expanded.to(inputs_embeds.device)
+
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+ n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
+ n_video_features = video_embeds.shape[0]
+ if n_video_tokens != n_video_features:
+ raise ValueError(
+ f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
+ )
+
+ mask = input_ids == self.config.video_token_id
+ mask_unsqueezed = mask.unsqueeze(-1)
+ mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
+ video_mask = mask_expanded.to(inputs_embeds.device)
+
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(inputs_embeds.device)
+
+ # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
+ if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
+ # calculate RoPE index once per generation in the pre-fill stage only
+ if (
+ (cache_position is not None and cache_position[0] == 0)
+ or self.rope_deltas is None
+ or (past_key_values is None or past_key_values.get_seq_length() == 0)
+ ):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ second_per_grid_ts,
+ attention_mask,
+ )
+ self.rope_deltas = rope_deltas
+ # then use the prev pre-calculated rope-deltas to get the correct position ids
+ else:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ delta = (
+ (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+ if cache_position is not None
+ else 0
+ )
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ if cache_position is not None: # otherwise `deltas` is an int `0`
+ delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ 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)
+
+ loss = None
+ if labels is not None:
+ # Upcast to float if we need to compute the loss to avoid potential precision issues
+ logits = logits.float()
+ # 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 Qwen2_5_VLCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ cache_position=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ second_per_grid_ts=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ second_per_grid_ts=second_per_grid_ts,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
+ model_inputs["position_ids"] = None
+
+ if cache_position[0] != 0:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+
+class Qwen2_5_VLVideosProcessorKwargs(VideosKwargs, total=False):
+ fps: Union[List[float], float]
+
+
+class Qwen2_5_VLImagesKwargs(Qwen2VLImagesKwargs):
+ pass
+
+
+class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: Qwen2_5_VLImagesKwargs
+ videos_kwargs: Qwen2_5_VLVideosProcessorKwargs
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ },
+ "videos_kwargs": {"fps": 2.0},
+ }
+
+
+class Qwen2_5_VLProcessor(Qwen2VLProcessor):
+ r"""
+ Constructs a Qwen2.5-VL processor which wraps a Qwen2.5-VL image processor and a Qwen2 tokenizer into a single processor.
+ [`Qwen2_5_VLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
+ [`~Qwen2_5_VLProcessor.__call__`] and [`~Qwen2_5_VLProcessor.decode`] for more information.
+ Args:
+ image_processor ([`Qwen2VLImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`Qwen2TokenizerFast`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ image_processor_class = "AutoImageProcessor"
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+ names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+ return names_from_processor + ["second_per_grid_ts"]
+
+ def __call__(
+ self,
+ images: ImageInput = None,
+ text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+ videos: VideoInput = None,
+ **kwargs: Unpack[Qwen2_5_VLProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+ and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+ the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
+ Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
+
+ Args:
+ 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.
+ 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).
+ videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+ tensor, or a nested list of 3D frames. 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:
+ [`BatchFeature`]: A [`BatchFeature`] 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`.
+ - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+ - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+ - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
+ - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ Qwen2_5_VLProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ if images is not None:
+ image_inputs = self.image_processor(images=images, videos=None, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+ else:
+ image_inputs = {}
+ image_grid_thw = None
+
+ if videos is not None:
+ videos_inputs = self.image_processor(images=None, videos=videos, **output_kwargs["images_kwargs"])
+ video_grid_thw = videos_inputs["video_grid_thw"]
+
+ fps = output_kwargs["videos_kwargs"].pop("fps", 2.0)
+ if isinstance(fps, (int, float)):
+ second_per_grid_ts = [self.image_processor.temporal_patch_size / fps] * len(video_grid_thw)
+ elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw):
+ second_per_grid_ts = [self.image_processor.temporal_patch_size / tmp for tmp in fps]
+ else:
+ raise ValueError(
+ f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number."
+ )
+ videos_inputs.update({"second_per_grid_ts": second_per_grid_ts})
+
+ else:
+ videos_inputs = {}
+ video_grid_thw = None
+
+ if not isinstance(text, list):
+ text = [text]
+
+ if image_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(text)):
+ while self.image_token in text[i]:
+ text[i] = text[i].replace(
+ self.image_token,
+ "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length),
+ 1,
+ )
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+ if video_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(text)):
+ while self.video_token in text[i]:
+ text[i] = text[i].replace(
+ self.video_token,
+ "<|placeholder|>" * (video_grid_thw[index].prod() // merge_length),
+ 1,
+ )
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.video_token)
+
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
+
+
+__all__ = [
+ "Qwen2_5_VLConfig",
+ "Qwen2_5_VLForConditionalGeneration",
+ "Qwen2_5_VLModel",
+ "Qwen2_5_VLPreTrainedModel",
+ "Qwen2_5_VLProcessor",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..e07642a1bf9d8242e0de7e517948cc1854727c63
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py
@@ -0,0 +1,239 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen2_5_vl.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen Team 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.
+from typing import List, Optional, Union
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, VideoInput
+from ...processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+
+
+class Qwen2_5_VLVideosProcessorKwargs(VideosKwargs, total=False):
+ fps: Union[List[float], float]
+
+
+class Qwen2_5_VLImagesKwargs(ImagesKwargs):
+ min_pixels: Optional[int]
+ max_pixels: Optional[int]
+ patch_size: Optional[int]
+ temporal_patch_size: Optional[int]
+ merge_size: Optional[int]
+
+
+class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: Qwen2_5_VLImagesKwargs
+ videos_kwargs: Qwen2_5_VLVideosProcessorKwargs
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ },
+ "videos_kwargs": {"fps": 2.0},
+ }
+
+
+class Qwen2_5_VLProcessor(ProcessorMixin):
+ r"""
+ Constructs a Qwen2.5-VL processor which wraps a Qwen2.5-VL image processor and a Qwen2 tokenizer into a single processor.
+ [`Qwen2_5_VLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
+ [`~Qwen2_5_VLProcessor.__call__`] and [`~Qwen2_5_VLProcessor.decode`] for more information.
+ Args:
+ image_processor ([`Qwen2VLImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`Qwen2TokenizerFast`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ attributes = ["image_processor", "tokenizer"]
+ valid_kwargs = ["chat_template"]
+
+ image_processor_class = "AutoImageProcessor"
+ tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
+
+ def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
+ self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+ self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ def __call__(
+ self,
+ images: ImageInput = None,
+ text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+ videos: VideoInput = None,
+ **kwargs: Unpack[Qwen2_5_VLProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+ and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+ the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
+ Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
+
+ Args:
+ 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.
+ 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).
+ videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+ tensor, or a nested list of 3D frames. 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:
+ [`BatchFeature`]: A [`BatchFeature`] 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`.
+ - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+ - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+ - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
+ - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ Qwen2_5_VLProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ if images is not None:
+ image_inputs = self.image_processor(images=images, videos=None, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+ else:
+ image_inputs = {}
+ image_grid_thw = None
+
+ if videos is not None:
+ videos_inputs = self.image_processor(images=None, videos=videos, **output_kwargs["images_kwargs"])
+ video_grid_thw = videos_inputs["video_grid_thw"]
+
+ fps = output_kwargs["videos_kwargs"].pop("fps", 2.0)
+ if isinstance(fps, (int, float)):
+ second_per_grid_ts = [self.image_processor.temporal_patch_size / fps] * len(video_grid_thw)
+ elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw):
+ second_per_grid_ts = [self.image_processor.temporal_patch_size / tmp for tmp in fps]
+ else:
+ raise ValueError(
+ f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number."
+ )
+ videos_inputs.update({"second_per_grid_ts": second_per_grid_ts})
+
+ else:
+ videos_inputs = {}
+ video_grid_thw = None
+
+ if not isinstance(text, list):
+ text = [text]
+
+ if image_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(text)):
+ while self.image_token in text[i]:
+ text[i] = text[i].replace(
+ self.image_token,
+ "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length),
+ 1,
+ )
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+ if video_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(text)):
+ while self.video_token in text[i]:
+ text[i] = text[i].replace(
+ self.video_token,
+ "<|placeholder|>" * (video_grid_thw[index].prod() // merge_length),
+ 1,
+ )
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.video_token)
+
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method forwards all its arguments to Qwen2TokenizerFast'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 Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+ the docstring of this method for more information.
+ """
+ return self.tokenizer.decode(*args, **kwargs)
+
+ def post_process_image_text_to_text(
+ self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+ ):
+ """
+ Post-process the output of the model to decode the text.
+
+ Args:
+ generated_outputs (`torch.Tensor` or `np.ndarray`):
+ The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+ or `(sequence_length,)`.
+ skip_special_tokens (`bool`, *optional*, defaults to `True`):
+ Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+ Clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+ Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+ **kwargs:
+ Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+ Returns:
+ `List[str]`: The decoded text.
+ """
+ return self.tokenizer.batch_decode(
+ generated_outputs,
+ skip_special_tokens=skip_special_tokens,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **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
+ names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+ return names_from_processor + ["second_per_grid_ts"]
+
+
+__all__ = ["Qwen2_5_VLProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__init__.py b/lib/python3.12/site-packages/transformers/models/rag/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a8f135ba454f08f1773239837a2db627ce075c6
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_rag import *
+ from .modeling_rag import *
+ from .modeling_tf_rag import *
+ from .retrieval_rag import *
+ from .tokenization_rag import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..118d6f306f66bba936cd7c18692ce9027a79887e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/configuration_rag.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/configuration_rag.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..205d8c55ce686049847eb4e3c776d3b4a0117338
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/configuration_rag.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_rag.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_rag.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b7855de5ef67baaaabc4b3b3d78528d4a2bbd4c3
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_rag.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_tf_rag.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_tf_rag.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f53a026e9a598b342940e5ae82894ab43f235d29
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/modeling_tf_rag.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/retrieval_rag.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/retrieval_rag.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3382be9f6951dd55c115f1856ddf4735d6b61ac9
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/retrieval_rag.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/__pycache__/tokenization_rag.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/tokenization_rag.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3a249cd56b06fb1c6ac6ae15463e93cf05b8b319
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/rag/__pycache__/tokenization_rag.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/rag/configuration_rag.py b/lib/python3.12/site-packages/transformers/models/rag/configuration_rag.py
new file mode 100644
index 0000000000000000000000000000000000000000..c76926f21879aa7ea5ab3d55a4be861db9491e97
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/configuration_rag.py
@@ -0,0 +1,186 @@
+# coding=utf-8
+# Copyright 2020, The RAG 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.
+"""RAG model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import add_start_docstrings
+
+
+RAG_CONFIG_DOC = r"""
+ [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and
+ can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ title_sep (`str`, *optional*, defaults to `" / "`):
+ Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].
+ doc_sep (`str`, *optional*, defaults to `" // "`):
+ Separator inserted between the text of the retrieved document and the original input when calling
+ [`RagRetriever`].
+ n_docs (`int`, *optional*, defaults to 5):
+ Number of documents to retrieve.
+ max_combined_length (`int`, *optional*, defaults to 300):
+ Max length of contextualized input returned by [`~RagRetriever.__call__`].
+ retrieval_vector_size (`int`, *optional*, defaults to 768):
+ Dimensionality of the document embeddings indexed by [`RagRetriever`].
+ retrieval_batch_size (`int`, *optional*, defaults to 8):
+ Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated
+ [`RagRetriever`].
+ dataset (`str`, *optional*, defaults to `"wiki_dpr"`):
+ A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids
+ using `datasets.list_datasets()`).
+ dataset_split (`str`, *optional*, defaults to `"train"`)
+ Which split of the `dataset` to load.
+ index_name (`str`, *optional*, defaults to `"compressed"`)
+ The index name of the index associated with the `dataset`. One can choose between `"legacy"`, `"exact"` and
+ `"compressed"`.
+ index_path (`str`, *optional*)
+ The path to the serialized faiss index on disk.
+ passages_path (`str`, *optional*):
+ A path to text passages compatible with the faiss index. Required if using
+ [`~models.rag.retrieval_rag.LegacyIndex`]
+ use_dummy_dataset (`bool`, *optional*, defaults to `False`)
+ Whether to load a "dummy" variant of the dataset specified by `dataset`.
+ label_smoothing (`float`, *optional*, defaults to 0.0):
+ Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing
+ in the loss calculation. If set to 0, no label smoothing is performed.
+ do_marginalize (`bool`, *optional*, defaults to `False`):
+ If `True`, the logits are marginalized over all documents by making use of
+ `torch.nn.functional.log_softmax`.
+ reduce_loss (`bool`, *optional*, defaults to `False`):
+ Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.
+ do_deduplication (`bool`, *optional*, defaults to `True`):
+ Whether or not to deduplicate the generations from different context documents for a given input. Has to be
+ set to `False` if used while training with distributed backend.
+ exclude_bos_score (`bool`, *optional*, defaults to `False`):
+ Whether or not to disregard the BOS token when computing the loss.
+ output_retrieved(`bool`, *optional*, defaults to `False`):
+ If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and
+ `context_attention_mask` are returned. See returned tensors for more detail.
+ 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*):
+ The id of the token to force as the last generated token when `max_length` is reached. Usually set to
+ `eos_token_id`.
+"""
+
+
+@add_start_docstrings(RAG_CONFIG_DOC)
+class RagConfig(PretrainedConfig):
+ model_type = "rag"
+ is_composition = True
+
+ def __init__(
+ self,
+ vocab_size=None,
+ is_encoder_decoder=True,
+ prefix=None,
+ bos_token_id=None,
+ pad_token_id=None,
+ eos_token_id=None,
+ decoder_start_token_id=None,
+ title_sep=" / ",
+ doc_sep=" // ",
+ n_docs=5,
+ max_combined_length=300,
+ retrieval_vector_size=768,
+ retrieval_batch_size=8,
+ dataset="wiki_dpr",
+ dataset_split="train",
+ index_name="compressed",
+ index_path=None,
+ passages_path=None,
+ use_dummy_dataset=False,
+ reduce_loss=False,
+ label_smoothing=0.0,
+ do_deduplication=True,
+ exclude_bos_score=False,
+ do_marginalize=False,
+ output_retrieved=False,
+ use_cache=True,
+ forced_eos_token_id=None,
+ dataset_revision=None,
+ **kwargs,
+ ):
+ super().__init__(
+ bos_token_id=bos_token_id,
+ pad_token_id=pad_token_id,
+ eos_token_id=eos_token_id,
+ decoder_start_token_id=decoder_start_token_id,
+ forced_eos_token_id=forced_eos_token_id,
+ is_encoder_decoder=is_encoder_decoder,
+ prefix=prefix,
+ vocab_size=vocab_size,
+ **kwargs,
+ )
+ if "question_encoder" not in kwargs or "generator" not in kwargs:
+ raise ValueError(
+ f"A configuraton of type {self.model_type} cannot be instantiated because "
+ f"both `question_encoder` and `generator` sub-configurations were not passed, only {kwargs}"
+ )
+ question_encoder_config = kwargs.pop("question_encoder")
+ question_encoder_model_type = question_encoder_config.pop("model_type")
+ decoder_config = kwargs.pop("generator")
+ decoder_model_type = decoder_config.pop("model_type")
+
+ from ..auto.configuration_auto import AutoConfig
+
+ self.question_encoder = AutoConfig.for_model(question_encoder_model_type, **question_encoder_config)
+ self.generator = AutoConfig.for_model(decoder_model_type, **decoder_config)
+
+ self.reduce_loss = reduce_loss
+ self.label_smoothing = label_smoothing
+ self.exclude_bos_score = exclude_bos_score
+ self.do_marginalize = do_marginalize
+
+ self.title_sep = title_sep
+ self.doc_sep = doc_sep
+ self.n_docs = n_docs
+ self.max_combined_length = max_combined_length
+
+ self.dataset = dataset
+ self.dataset_split = dataset_split
+ self.index_name = index_name
+
+ self.retrieval_vector_size = retrieval_vector_size
+ self.retrieval_batch_size = retrieval_batch_size
+ self.passages_path = passages_path
+ self.index_path = index_path
+ self.use_dummy_dataset = use_dummy_dataset
+ self.dataset_revision = dataset_revision
+
+ self.output_retrieved = output_retrieved
+
+ self.do_deduplication = do_deduplication
+
+ self.use_cache = use_cache
+
+ if self.forced_eos_token_id is None:
+ self.forced_eos_token_id = getattr(self.generator, "forced_eos_token_id", None)
+
+ @classmethod
+ def from_question_encoder_generator_configs(
+ cls, question_encoder_config: PretrainedConfig, generator_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
+ """
+ return cls(question_encoder=question_encoder_config.to_dict(), generator=generator_config.to_dict(), **kwargs)
+
+
+__all__ = ["RagConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/rag/modeling_rag.py b/lib/python3.12/site-packages/transformers/models/rag/modeling_rag.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2258d9767a84306ff37cbd07a02e0bbe1c23e7b
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/modeling_rag.py
@@ -0,0 +1,1627 @@
+# coding=utf-8
+# Copyright 2020, The RAG 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.
+"""RAG model implementation."""
+
+import copy
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from ...configuration_utils import PretrainedConfig
+from ...generation import GenerationConfig, LogitsProcessorList, StoppingCriteriaList
+from ...modeling_outputs import ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_rag import RagConfig
+from .retrieval_rag import RagRetriever
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "RagConfig"
+
+
+@dataclass
+class RetrievAugLMMarginOutput(ModelOutput):
+ """
+ Base class for retriever augmented marginalized models outputs.
+
+ 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. The score is possibly marginalized over all documents for
+ each vocabulary token.
+ doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+ past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size,
+ num_heads, sequence_length, embed_size_per_head)`).
+
+ Contains precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used
+ (see `past_key_values` input) to speed up sequential decoding.
+ retrieved_doc_embeds (`torch.FloatTensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*):
+ Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute
+ the `doc_scores`.
+ retrieved_doc_ids (`torch.LongTensor` of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*):
+ The indexes of the embedded documents retrieved by the retriever.
+ context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input ids post-processed from the retrieved documents and the question encoder input_ids by the retriever.
+ context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+ question_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 question encoder pooled output of the
+ model.
+ question_enc_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the question encoder at the output of each layer plus the initial embedding outputs.
+ question_enc_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 question encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_enc_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 generator encoder of the model.
+ generator_enc_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs.
+ generator_enc_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 generator encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_dec_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs.
+ generator_dec_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 generator decoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_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)`.
+
+ Cross-attentions weights of the generator decoder, after the attention softmax, used to compute the
+ weighted average in the cross-attention heads.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ doc_scores: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[List[torch.FloatTensor]] = None
+ retrieved_doc_embeds: Optional[torch.FloatTensor] = None
+ retrieved_doc_ids: Optional[torch.LongTensor] = None
+ context_input_ids: Optional[torch.LongTensor] = None
+ context_attention_mask: Optional[torch.LongTensor] = None
+ question_encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ question_enc_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ question_enc_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_enc_last_hidden_state: Optional[torch.FloatTensor] = None
+ generator_enc_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_enc_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_dec_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_dec_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class RetrievAugLMOutput(ModelOutput):
+ """
+ Args:
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head. The score is possibly marginalized over all documents for
+ each vocabulary token.
+ doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+ past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size,
+ num_heads, sequence_length, embed_size_per_head)`).
+
+ Contains precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used
+ (see `past_key_values` input) to speed up sequential decoding.
+ retrieved_doc_embeds (`torch.FloatTensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*):
+ Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute
+ the `doc_scores`.
+ retrieved_doc_ids (`torch.LongTensor` of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*):
+ The indexes of the embedded documents retrieved by the retriever.
+ context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input ids post-processed from the retrieved documents and the question encoder input_ids by the retriever.
+ context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+ question_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 question encoder pooled output of the
+ model.
+ question_enc_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the question encoder at the output of each layer plus the initial embedding outputs.
+ question_enc_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 question encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_enc_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 generator encoder of the model.
+ generator_enc_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs.
+ generator_enc_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 generator encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_dec_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 and one for the output of each layer) of
+ shape `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs.
+ generator_dec_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 generator decoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_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)`.
+
+ Cross-attentions weights of the generator decoder, after the attention softmax, used to compute the
+ weighted average in the cross-attention heads.
+ """
+
+ logits: Optional[torch.FloatTensor] = None
+ doc_scores: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[List[torch.FloatTensor]] = None
+ retrieved_doc_embeds: Optional[torch.FloatTensor] = None
+ retrieved_doc_ids: Optional[torch.LongTensor] = None
+ context_input_ids: Optional[torch.LongTensor] = None
+ context_attention_mask: Optional[torch.LongTensor] = None
+ question_encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+ question_enc_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ question_enc_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_enc_last_hidden_state: Optional[torch.FloatTensor] = None
+ generator_enc_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_enc_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_dec_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_dec_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+ generator_cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+class RagPreTrainedModel(PreTrainedModel):
+ r"""
+ RAG models were released with the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP
+ Tasks](https://arxiv.org/abs/2005.11401) by Patrick Lewis, Ethan Perez, Aleksandra Piktus et al.
+
+ RAG is a retriever augmented model and encapsulate three components: a question encoder, a dataset retriever and a
+ generator, the encoder and generator are trainable while the retriever is just an indexed dataset.
+
+ """
+
+ config_class = RagConfig
+ base_model_prefix = "rag"
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ @classmethod
+ def from_pretrained_question_encoder_generator(
+ cls,
+ question_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ generator_pretrained_model_name_or_path: Optional[str] = None,
+ retriever: RagRetriever = None,
+ **kwargs,
+ ) -> PreTrainedModel:
+ r"""
+ Instantiates an question encoder and a generator 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:
+ question_encoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the question 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.
+
+ generator_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the generator. 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.
+ retriever ([`RagRetriever`], *optional*):
+ The retriever to use.
+ kwwargs (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 question_encoder configuration, use the prefix *question_encoder_* for each
+ configuration parameter.
+ - To update the generator configuration, use the prefix *generator_* 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 RagModel
+
+ >>> # initialize a RAG from two pretrained models.
+ >>> model = RagModel.from_pretrained_question_encoder_generator(
+ ... "facebook/dpr-question_encoder-single-nq-base", "google-t5/t5-small"
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./rag")
+ >>> # load fine-tuned model
+ >>> model = RagModel.from_pretrained("./rag")
+ ```"""
+
+ kwargs_question_encoder = {
+ argument[len("question_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("question_encoder_")
+ }
+
+ kwargs_generator = {
+ argument[len("generator_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("generator_")
+ }
+
+ # remove question_encoder, generator kwargs from kwargs
+ for key in kwargs_question_encoder.keys():
+ del kwargs["question_encoder_" + key]
+ for key in kwargs_generator.keys():
+ del kwargs["generator_" + key]
+
+ # Load and initialize the question_encoder and generator
+ # The distinction between question_encoder and generator at the model level is made
+ # by the value of the flag `is_generator` that we need to set correctly.
+ question_encoder = kwargs_question_encoder.pop("model", None)
+ if question_encoder is None:
+ assert question_encoder_pretrained_model_name_or_path is not None, (
+ "If `model` is not defined as an argument, a `question_encoder_pretrained_model_name_or_path` has to"
+ " be defined"
+ )
+ from ..auto.modeling_auto import AutoModel
+
+ if "config" not in kwargs_question_encoder:
+ from ..auto.configuration_auto import AutoConfig
+
+ question_encoder_config, kwargs_question_encoder = AutoConfig.from_pretrained(
+ question_encoder_pretrained_model_name_or_path,
+ **kwargs_question_encoder,
+ return_unused_kwargs=True,
+ )
+ kwargs_question_encoder["config"] = question_encoder_config
+
+ question_encoder = AutoModel.from_pretrained(
+ question_encoder_pretrained_model_name_or_path, **kwargs_question_encoder
+ )
+
+ generator = kwargs_generator.pop("model", None)
+ if generator is None:
+ assert generator_pretrained_model_name_or_path is not None, (
+ "If `generator_model` is not defined as an argument, a `generator_pretrained_model_name_or_path` has"
+ " to be defined"
+ )
+ from ..auto.modeling_auto import AutoModelForSeq2SeqLM
+
+ if "config" not in kwargs_generator:
+ from ..auto.configuration_auto import AutoConfig
+
+ generator_config, kwargs_generator = AutoConfig.from_pretrained(
+ generator_pretrained_model_name_or_path, **kwargs_generator, return_unused_kwargs=True
+ )
+
+ kwargs_generator["config"] = generator_config
+
+ generator = AutoModelForSeq2SeqLM.from_pretrained(
+ generator_pretrained_model_name_or_path, **kwargs_generator
+ )
+
+ # instantiate config with corresponding kwargs
+ config = kwargs.get("config", None)
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+
+ return cls(question_encoder=question_encoder, generator=generator, config=config, retriever=retriever)
+
+
+RAG_START_DOCSTRING = r"""
+
+ RAG is a seq2seq model which encapsulates two core components: a question encoder and a generator. During a forward
+ pass, we encode the input with the question encoder and pass it to the retriever to extract relevant context
+ documents. The documents are then prepended to the input. Such contextualized inputs is passed to the generator.
+
+ The question encoder can be any *autoencoding* model, preferably [`DPRQuestionEncoder`], and the generator can be
+ any *seq2seq* model, preferably [`BartForConditionalGeneration`].
+
+ The model can be initialized with a [`RagRetriever`] for end-to-end generation or used in combination with the
+ outputs of a retriever in multiple steps---see examples for more details. The model is compatible any
+ *autoencoding* model as the `question_encoder` and any *seq2seq* model with language model head as the `generator`.
+ It has been tested with [`DPRQuestionEncoder`] as the `question_encoder` and [`BartForConditionalGeneration`] or
+ [`T5ForConditionalGeneration`] as the `generator`.
+
+ 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.
+
+
+ Args:
+ config ([`RagConfig`]):
+ 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.
+ question_encoder ([`PreTrainedModel`]):
+ An encoder model compatible with the faiss index encapsulated by the `retriever`.
+ generator ([`PreTrainedModel`]):
+ A seq2seq model used as the generator in the RAG architecture.
+ retriever ([`RagRetriever`]):
+ A retriever class encapsulating a faiss index queried to obtain context documents for current inputs.
+"""
+
+
+RAG_FORWARD_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies
+ which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to
+ obtain the indices.
+
+ [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)
+ encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*)
+ Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`,
+ *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs *
+ sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the
+ generator's encoder.
+
+ Used by the ([`RagModel`]) model during decoding.
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Provide for generation tasks. `None` by default, construct as per instructions for the generator model
+ you're using with your RAG instance.
+ 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.
+ past_key_values (`tuple(tuple(torch.FloatTensor))`):
+ Tuple consists of two elements: `encoder_outputs` of the RAG model (see `encoder_outputs`) and
+ `past_key_values` of the underlying generator. Can be used to speed up decoding. `past_key_values` are used
+ in the ([`RagTokenForGeneration`]) model during decoding.
+ doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`. If the model has is not initialized with a `retriever` `doc_scores`
+ has to be provided to the forward pass. `doc_scores` can be computed via
+ `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information.
+ context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever. If the model was not initialized with a `retriever` ``context_input_ids` has to be provided to
+ the forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`,*optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever. If the model has is not initialized with a `retriever` `context_attention_mask` has to be
+ provided to the forward pass. `context_attention_mask` are returned by [`~RagRetriever.__call__`].
+ 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.
+ output_retrieved(`bool`, *optional*):
+ Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and
+ `context_attention_mask`. See returned tensors for more detail.
+ n_docs (`int`, *optional*, defaults to `config.n_docs``)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+"""
+
+
+@add_start_docstrings_to_model_forward(RAG_START_DOCSTRING)
+class RagModel(RagPreTrainedModel):
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[PreTrainedModel] = None,
+ generator: Optional[PreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None, # or maybe just use a `set_retriever(...)` method
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an question_encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+ else:
+ assert isinstance(config, self.config_class), f"config: {config} has to be of type {self.config_class}"
+ super().__init__(config)
+ if question_encoder is None:
+ from ..auto.modeling_auto import AutoModel
+
+ question_encoder = AutoModel.from_config(config.question_encoder)
+
+ if generator is None:
+ from ..auto.modeling_auto import AutoModelForSeq2SeqLM
+
+ generator = AutoModelForSeq2SeqLM.from_config(config.generator)
+
+ self.retriever = retriever
+ if self.retriever is not None:
+ assert isinstance(retriever, RagRetriever), (
+ f"`self.retriever` is of type {type(self.retriever)}, but should be of type `RagRetriever`"
+ )
+ self.retriever = retriever
+
+ self.question_encoder = question_encoder
+ self.generator = generator
+
+ self.ctx_encoder = None
+ self.context_encoder_training = False
+
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=RetrievAugLMOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ decoder_input_ids: Optional[torch.LongTensor] = None,
+ decoder_attention_mask: Optional[torch.BoolTensor] = None,
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ doc_scores: Optional[torch.FloatTensor] = None,
+ context_input_ids: Optional[torch.LongTensor] = None,
+ context_attention_mask: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_retrieved: Optional[bool] = None,
+ n_docs: Optional[int] = None,
+ ) -> Union[Tuple[torch.Tensor], RetrievAugLMOutput]:
+ r"""
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RagRetriever, RagModel
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = RagModel.from_pretrained("facebook/rag-token-base", retriever=retriever)
+
+ >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
+ >>> outputs = model(input_ids=inputs["input_ids"])
+ ```"""
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ 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_retrieved = output_retrieved if output_retrieved is not None else self.config.output_retrieved
+
+ # whether retriever has to be used
+ has_to_retrieve = (
+ self.retriever is not None
+ and (context_input_ids is None or context_attention_mask is None or doc_scores is None)
+ and encoder_outputs is None
+ )
+ # encoder_outputs are pre-computed during RAG-token generation
+ if encoder_outputs is None:
+ if has_to_retrieve:
+ question_enc_outputs = self.question_encoder(
+ input_ids, attention_mask=attention_mask, return_dict=True
+ )
+ question_encoder_last_hidden_state = question_enc_outputs[0] # hidden states of question encoder
+
+ retriever_outputs = self.retriever(
+ input_ids,
+ question_encoder_last_hidden_state.detach().to(device="cpu", dtype=torch.float32).numpy(),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="pt",
+ )
+ if self.context_encoder_training:
+ (
+ context_input_ids,
+ context_attention_mask,
+ retrieved_doc_embeds,
+ retrived_doc_input_ids,
+ retrived_doc_attention_mask,
+ retrieved_doc_ids,
+ ) = (
+ retriever_outputs["context_input_ids"],
+ retriever_outputs["context_attention_mask"],
+ retriever_outputs["retrieved_doc_embeds"],
+ retriever_outputs["tokenized_doc_ids"],
+ retriever_outputs["tokenized_doc_attention_mask"],
+ retriever_outputs["doc_ids"],
+ )
+
+ context_input_ids = context_input_ids.to(input_ids)
+ context_attention_mask = context_attention_mask.to(input_ids)
+
+ retrived_doc_input_ids = retrived_doc_input_ids.to(input_ids)
+ retrived_doc_attention_mask = retrived_doc_attention_mask.to(input_ids)
+ retrieved_doc_embeds = self.ctx_encoder(
+ retrived_doc_input_ids, attention_mask=retrived_doc_attention_mask, return_dict=True
+ ).pooler_output
+ retrieved_doc_embeds = retrieved_doc_embeds.view(
+ -1, n_docs, question_encoder_last_hidden_state.shape[1]
+ ) # reshaping
+
+ # compute doc_scores involving ctx_encoder
+ doc_scores = torch.bmm(
+ question_encoder_last_hidden_state.unsqueeze(1), retrieved_doc_embeds.transpose(1, 2)
+ ).squeeze(1)
+
+ else:
+ context_input_ids, context_attention_mask, retrieved_doc_embeds, retrieved_doc_ids = (
+ retriever_outputs["context_input_ids"],
+ retriever_outputs["context_attention_mask"],
+ retriever_outputs["retrieved_doc_embeds"],
+ retriever_outputs["doc_ids"],
+ )
+
+ # set to correct device
+ retrieved_doc_embeds = retrieved_doc_embeds.to(question_encoder_last_hidden_state)
+ context_input_ids = context_input_ids.to(input_ids)
+ context_attention_mask = context_attention_mask.to(input_ids)
+
+ # compute doc_scores
+ doc_scores = torch.bmm(
+ question_encoder_last_hidden_state.unsqueeze(1), retrieved_doc_embeds.transpose(1, 2)
+ ).squeeze(1)
+ else:
+ assert context_input_ids is not None, (
+ "Make sure that `context_input_ids` are passed, if no `retriever` is set. Alternatively, you can"
+ " set a retriever using the `set_retriever(...)` function."
+ )
+ assert context_attention_mask is not None, (
+ "Make sure that `context_attention_mask` are passed, if no `retriever` is set. Alternatively, you"
+ " can set a retriever using the `set_retriever(...)` function."
+ )
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed, if no `retriever` is set. Alternatively, you can set a"
+ " retriever using the `set_retriever(...)` function."
+ )
+
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed when passing `encoder_outputs` to the forward function."
+ )
+
+ assert (doc_scores.shape[1] % n_docs) == 0, (
+ f" The first dimension of `context_input_ids` should be a multiple of `n_docs`={n_docs}, but is"
+ f" {context_input_ids.shape[0]}."
+ )
+
+ # Decoder input without context documents
+ if decoder_input_ids is not None:
+ decoder_input_ids = decoder_input_ids.repeat_interleave(n_docs, dim=0)
+
+ if decoder_attention_mask is not None:
+ decoder_attention_mask = decoder_attention_mask.repeat_interleave(n_docs, dim=0)
+
+ gen_outputs = self.generator(
+ input_ids=context_input_ids,
+ attention_mask=context_attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ return_dict=True,
+ )
+
+ if not has_to_retrieve:
+ question_encoder_last_hidden_state = None
+ question_enc_hidden_states = None
+ question_enc_attentions = None
+ retrieved_doc_embeds = None
+ retrieved_doc_ids = None
+ else:
+ question_enc_hidden_states = question_enc_outputs.hidden_states
+ question_enc_attentions = question_enc_outputs.attentions
+
+ if not has_to_retrieve or not output_retrieved:
+ # don't output retrieved docs
+ context_input_ids = (None,)
+ context_attention_mask = None
+ retrieved_doc_embeds = None
+ retrieved_doc_ids = None
+
+ return RetrievAugLMOutput(
+ logits=gen_outputs.logits,
+ doc_scores=doc_scores,
+ past_key_values=gen_outputs.past_key_values,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ retrieved_doc_embeds=retrieved_doc_embeds,
+ retrieved_doc_ids=retrieved_doc_ids,
+ question_encoder_last_hidden_state=question_encoder_last_hidden_state,
+ question_enc_hidden_states=question_enc_hidden_states,
+ question_enc_attentions=question_enc_attentions,
+ generator_enc_last_hidden_state=gen_outputs.encoder_last_hidden_state,
+ generator_enc_hidden_states=gen_outputs.encoder_hidden_states,
+ generator_enc_attentions=gen_outputs.encoder_attentions,
+ generator_dec_hidden_states=gen_outputs.decoder_hidden_states,
+ generator_dec_attentions=gen_outputs.decoder_attentions,
+ generator_cross_attentions=gen_outputs.cross_attentions,
+ )
+
+
+@add_start_docstrings_to_model_forward(
+ """
+ A RAG-sequence model implementation. It performs RAG-sequence specific marginalization in the forward pass.
+ """,
+ RAG_START_DOCSTRING,
+)
+class RagSequenceForGeneration(RagPreTrainedModel):
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[PreTrainedModel] = None,
+ generator: Optional[PreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None,
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+ super().__init__(config)
+
+ # instantiate model
+ self.rag = RagModel(config=config, question_encoder=question_encoder, generator=generator, retriever=retriever)
+
+ def set_retriever(self, retriever: RagRetriever):
+ self.rag.retriever = retriever
+
+ def set_context_encoder_for_training(self, ctx_encoder: PreTrainedModel):
+ self.rag.context_encoder_training = True
+ self.rag.ctx_encoder = ctx_encoder
+
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=RetrievAugLMMarginOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+ decoder_input_ids: Optional[torch.LongTensor] = None,
+ decoder_attention_mask: Optional[torch.BoolTensor] = None,
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+ context_input_ids: Optional[torch.LongTensor] = None,
+ context_attention_mask: Optional[torch.LongTensor] = None,
+ doc_scores: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_retrieved: Optional[bool] = None,
+ exclude_bos_score: Optional[bool] = None,
+ reduce_loss: Optional[bool] = None,
+ labels: Optional[torch.LongTensor] = None,
+ n_docs: Optional[int] = None,
+ **kwargs, # needs kwargs for generation
+ ) -> RetrievAugLMMarginOutput:
+ r"""
+ exclude_bos_score (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the score of the BOS token is disregarded when computing
+ the loss.
+ reduce_loss (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `torch.Tensor.sum`
+ operation.
+ kwargs (`Dict[str, any]`, *optional*, defaults to `{}`):
+ Legacy dictionary, which is required so that model can use *generate()* function.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RagRetriever, RagSequenceForGeneration
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
+
+ >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
+ >>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt")
+ >>> input_ids = inputs["input_ids"]
+ >>> labels = targets["input_ids"]
+ >>> outputs = model(input_ids=input_ids, labels=labels)
+
+ >>> # or use retriever separately
+ >>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True)
+ >>> # 1. Encode
+ >>> question_hidden_states = model.question_encoder(input_ids)[0]
+ >>> # 2. Retrieve
+ >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
+ >>> doc_scores = torch.bmm(
+ ... question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)
+ ... ).squeeze(1)
+ >>> # 3. Forward to generator
+ >>> outputs = model(
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... decoder_input_ids=labels,
+ ... )
+ ```"""
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ exclude_bos_score = exclude_bos_score if exclude_bos_score is not None else self.config.exclude_bos_score
+ reduce_loss = reduce_loss if reduce_loss is not None else self.config.reduce_loss
+
+ if labels is not None:
+ if decoder_input_ids is None:
+ decoder_input_ids = labels
+ use_cache = False
+
+ outputs = self.rag(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ doc_scores=doc_scores,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ output_retrieved=output_retrieved,
+ n_docs=n_docs,
+ )
+
+ loss = None
+ if labels is not None:
+ loss = self.get_nll(
+ outputs.logits,
+ outputs.doc_scores,
+ decoder_input_ids,
+ reduce_loss=reduce_loss,
+ epsilon=self.config.label_smoothing,
+ exclude_bos_score=exclude_bos_score,
+ n_docs=n_docs,
+ )
+
+ return RetrievAugLMMarginOutput(
+ loss=loss,
+ logits=outputs.logits,
+ doc_scores=outputs.doc_scores,
+ past_key_values=outputs.past_key_values,
+ context_input_ids=outputs.context_input_ids,
+ context_attention_mask=outputs.context_attention_mask,
+ retrieved_doc_embeds=outputs.retrieved_doc_embeds,
+ retrieved_doc_ids=outputs.retrieved_doc_ids,
+ question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state,
+ question_enc_hidden_states=outputs.question_enc_hidden_states,
+ question_enc_attentions=outputs.question_enc_attentions,
+ generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state,
+ generator_enc_hidden_states=outputs.generator_enc_hidden_states,
+ generator_enc_attentions=outputs.generator_enc_attentions,
+ generator_dec_hidden_states=outputs.generator_dec_hidden_states,
+ generator_dec_attentions=outputs.generator_dec_attentions,
+ generator_cross_attentions=outputs.generator_cross_attentions,
+ )
+
+ @property
+ def retriever(self):
+ return self.rag.retriever
+
+ @property
+ def generator(self):
+ return self.rag.generator
+
+ @property
+ def question_encoder(self):
+ return self.rag.question_encoder
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.LongTensor] = None,
+ context_input_ids: Optional[torch.LongTensor] = None,
+ context_attention_mask: Optional[torch.LongTensor] = None,
+ doc_scores: Optional[torch.FloatTensor] = None,
+ do_deduplication: Optional[bool] = None, # defaults to True
+ num_return_sequences: Optional[int] = None, # defaults to 1
+ num_beams: Optional[int] = None, # defaults to 1
+ n_docs: Optional[int] = None,
+ **model_kwargs,
+ ) -> torch.LongTensor:
+ """
+ Implements RAG sequence "thorough" decoding. Read the [`~generation.GenerationMixin.generate`]` documentation
+ for more information on how to set other generate input parameters.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ The sequence used as a prompt for the generation. If `input_ids` is not passed, then
+ `context_input_ids` has to be provided.
+ 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)
+ context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder input_ids by the
+ retriever.
+ context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model is not initialized with a `retriever` or `input_ids` is not given, `context_input_ids` and
+ `context_attention_mask` have to be provided to the forward pass. They are returned by
+ [`~RagRetriever.__call__`].
+ doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+
+ If the model is not initialized with a `retriever` or `input_ids` is not given, `doc_scores` has to be
+ provided to the forward pass. `doc_scores` are returned by [`~RagRetriever.__call__`].
+ do_deduplication (`bool`, *optional*):
+ Whether or not to deduplicate the generations from different context documents for a given input. Has
+ to be set to `False` if used while training with distributed backend.
+ num_return_sequences(`int`, *optional*, defaults to 1):
+ The number of independently computed returned sequences for each element in the batch. Note that this
+ is not the value we pass to the `generator`'s `[`~generation.GenerationMixin.generate`]` function,
+ where we set `num_return_sequences` to `num_beams`.
+ num_beams (`int`, *optional*, defaults to 1):
+ Number of beams for beam search. 1 means no beam search.
+ n_docs (`int`, *optional*, defaults to `config.n_docs`)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional kwargs will be passed to [`~generation.GenerationMixin.generate`].
+
+ Return:
+ `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated
+ sequences. The second dimension (sequence length) is either equal to `max_length` or shorter if all batches
+ finished early due to the `eos_token_id`.
+ """
+
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ do_deduplication = do_deduplication if do_deduplication is not None else self.config.do_deduplication
+ num_doc_return_sequences = (
+ num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
+ )
+ num_beams = num_beams if num_beams is not None else self.config.num_beams
+
+ assert input_ids is not None or context_input_ids is not None, (
+ " At least one of input_ids or context_input_ids must be given"
+ )
+
+ if self.retriever is not None and context_input_ids is None:
+ question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
+ context_input_ids = self.retriever(
+ input_ids,
+ question_hidden_states.detach().to(device="cpu", dtype=torch.float32).numpy(),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="pt",
+ )["context_input_ids"]
+
+ # set to correct device
+ context_input_ids = context_input_ids.to(input_ids)
+
+ hypos = []
+ model_kwargs["num_beams"] = num_beams
+ model_kwargs["num_return_sequences"] = num_beams
+ model_kwargs["attention_mask"] = None
+
+ batch_size = input_ids.shape[0] if input_ids is not None else context_input_ids.shape[0] // n_docs
+
+ for index in range(batch_size):
+ # first, generate beams from documents:
+ generator_input_ids = context_input_ids[index * n_docs : (index + 1) * n_docs] # (n_docs, max_len)
+
+ output_sequences = self.generator.generate(
+ generator_input_ids,
+ **model_kwargs,
+ ) # n_docs * n_beam, tgt_len
+ if do_deduplication:
+ # do_deduplication, max_output_len
+ output_sequences = torch.stack(list({str(k.tolist()): k for k in output_sequences}.values()))
+
+ num_candidates = output_sequences.shape[
+ 0
+ ] # after deduplication, this number can be less than n_docs*n_beam
+
+ # then, run model forwards to get nll scores:
+ if input_ids is not None:
+ new_input_ids = input_ids[index : index + 1].repeat(num_candidates, 1)
+ outputs = self(new_input_ids, labels=output_sequences, exclude_bos_score=True)
+ else: # input_ids is None, need context_input_ids/mask and doc_scores
+ assert context_attention_mask is not None, (
+ "Make sure that `context_attention_mask` are passed, if no `input_ids` is set. Alternatively, you"
+ " can set a retriever using the `set_retriever(...)` function."
+ )
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed, if no `input_ids` is set. Alternatively, you can set a"
+ " retriever using the `set_retriever(...)` function."
+ )
+
+ individual_input_ids = generator_input_ids.repeat(
+ num_candidates, 1
+ ) # (num_candidates*n_docs, max_len)
+
+ individual_attention_mask = context_attention_mask[index * n_docs : (index + 1) * n_docs]
+ individual_attention_mask = individual_attention_mask.repeat(num_candidates, 1)
+
+ individual_doc_scores = doc_scores[index : (index + 1), :] # doc_scores.shape = [batch, n_docs]
+ individual_doc_scores = individual_doc_scores.repeat(num_candidates, 1) # [num_candidates, n_docs]
+
+ outputs = self(
+ context_input_ids=individual_input_ids,
+ context_attention_mask=individual_attention_mask,
+ doc_scores=individual_doc_scores,
+ labels=output_sequences,
+ exclude_bos_score=True,
+ )
+
+ top_cand_inds = (-outputs["loss"]).topk(num_doc_return_sequences)[1]
+
+ # add hypothesis
+ hypos.append(output_sequences[top_cand_inds])
+
+ return self._cat_and_pad(hypos, pad_token_id=self.config.generator.pad_token_id)
+
+ def get_nll(
+ self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, exclude_bos_score=False, n_docs=None
+ ):
+ # shift tokens left
+ target = torch.cat(
+ [target[:, 1:], target.new(target.shape[0], 1).fill_(self.config.generator.pad_token_id)], 1
+ )
+
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # bos_token_id is None for T5
+ bos_token_id = self.config.bos_token_id or self.config.generator.bos_token_id
+ use_bos = bos_token_id is not None and target[:, 0].eq(bos_token_id).all()
+
+ def _mask_pads(ll, smooth_obj):
+ pad_mask = target.eq(self.config.generator.pad_token_id)
+ if pad_mask.any():
+ ll.masked_fill_(pad_mask, 0.0)
+ smooth_obj.masked_fill_(pad_mask, 0.0)
+ return ll.squeeze(-1), smooth_obj.squeeze(-1)
+
+ # seq_logits dim = (batch*n_docs, tgt_len , #vocabs)
+ seq_logprobs = nn.functional.log_softmax(seq_logits, dim=-1).view(
+ seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.size(-1)
+ ) # batch_size x n_docs x tgt_len x #vocab_size
+ doc_logprobs = nn.functional.log_softmax(doc_scores, dim=1).unsqueeze(-1).unsqueeze(-1)
+
+ # RAG-sequence marginalization
+ first_token_scores = seq_logprobs[:, :, :1, :]
+ second_token_scores = seq_logprobs[:, :, 1:2, :]
+ remainder = seq_logprobs[:, :, 2:, :]
+ rag_logprobs = torch.cat([first_token_scores, second_token_scores + doc_logprobs, remainder], dim=2)
+
+ # calculate loss
+ target = target.unsqueeze(1).unsqueeze(-1).repeat(1, n_docs, 1, 1)
+ assert target.dim() == rag_logprobs.dim()
+
+ ll = rag_logprobs.gather(dim=-1, index=target)
+ smooth_obj = rag_logprobs.sum(dim=-1, keepdim=True) # total sum of all (normalised) logits
+
+ ll, smooth_obj = _mask_pads(ll, smooth_obj)
+
+ # sum over tokens, exclude bos while scoring
+ ll = ll[:, :, 1:].sum(2) if exclude_bos_score and use_bos else ll.sum(2)
+ smooth_obj = smooth_obj.sum(2)
+ ll = ll.logsumexp(1) # logsumexp over docs
+ smooth_obj = smooth_obj.logsumexp(1)
+
+ nll_loss = -ll
+ smooth_loss = -smooth_obj
+
+ if reduce_loss:
+ nll_loss = nll_loss.sum()
+ smooth_loss = smooth_loss.sum()
+
+ eps_i = epsilon / rag_logprobs.size(-1)
+ loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
+ return loss
+
+ @staticmethod
+ def _cat_and_pad(tensors, pad_token_id):
+ output = (
+ tensors[0].new(sum([t.shape[0] for t in tensors]), max([t.shape[1] for t in tensors])).fill_(pad_token_id)
+ )
+ ind = 0
+ for t in tensors:
+ output[ind : ind + t.shape[0], : t.shape[1]] = t
+ ind += t.shape[0]
+ return output
+
+
+@add_start_docstrings_to_model_forward(
+ """
+ A RAG-token model implementation. It performs RAG-token specific marginalization in the forward pass.
+ """,
+ RAG_START_DOCSTRING,
+)
+class RagTokenForGeneration(RagPreTrainedModel):
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[PreTrainedModel] = None,
+ generator: Optional[PreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None,
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+
+ super().__init__(config)
+
+ # instantiate model
+ self.rag = RagModel(config=config, question_encoder=question_encoder, generator=generator, retriever=retriever)
+
+ def set_retriever(self, retriever: RagRetriever):
+ self.rag.retriever = retriever
+
+ def set_context_encoder_for_training(self, ctx_encoder: PreTrainedModel):
+ self.rag.context_encoder_training = True
+ self.rag.ctx_encoder = ctx_encoder
+
+ def prepare_inputs_for_generation(
+ self,
+ decoder_input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ use_cache=None,
+ encoder_outputs=None,
+ doc_scores=None,
+ n_docs=None,
+ **kwargs,
+ ):
+ # Overwritten -- `do_marginalize` is explicitly set in the output
+
+ if past_key_values is not None:
+ # if past is defined use only last decoder_input_ids
+ decoder_input_ids = decoder_input_ids[:, -1:]
+
+ return {
+ "input_ids": None,
+ "encoder_outputs": encoder_outputs,
+ "doc_scores": doc_scores,
+ "context_attention_mask": attention_mask,
+ "decoder_input_ids": decoder_input_ids,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ "do_marginalize": True,
+ "n_docs": n_docs,
+ }
+
+ @property
+ def retriever(self):
+ return self.rag.retriever
+
+ @property
+ def generator(self):
+ return self.rag.generator
+
+ @property
+ def question_encoder(self):
+ return self.rag.question_encoder
+
+ @staticmethod
+ def _reorder_cache(past_key_values, beam_idx):
+ """Reorders cache for generation. BART-inspired but we need to take care of the extra dimension for docs"""
+
+ def _reorder_stacked(hidden_states, new_order):
+ n_docs = hidden_states.shape[0] // new_order.shape[0]
+ hidden_states = hidden_states.view(-1, n_docs, *hidden_states.shape[1:])
+ hidden_states = hidden_states.index_select(0, new_order)
+ result = hidden_states.view(-1, *hidden_states.shape[2:])
+ return result
+
+ 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
+ reordered_past += (
+ tuple(_reorder_stacked(past_state, beam_idx.to(past_state.device)) for past_state in layer_past),
+ )
+
+ return reordered_past
+
+ def marginalize(self, seq_logits, doc_scores, n_docs=None):
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # RAG-token marginalization
+ seq_logprobs = nn.functional.log_softmax(seq_logits, dim=-1).view(
+ seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.size(-1)
+ )
+ doc_logprobs = torch.log_softmax(doc_scores, dim=1)
+ log_prob_sum = seq_logprobs + doc_logprobs.unsqueeze(-1).unsqueeze(-1)
+ return torch.logsumexp(log_prob_sum, dim=1)
+
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=RetrievAugLMMarginOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+ decoder_input_ids: Optional[torch.LongTensor] = None,
+ decoder_attention_mask: Optional[torch.BoolTensor] = None,
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+ context_input_ids: Optional[torch.LongTensor] = None,
+ context_attention_mask: Optional[torch.LongTensor] = None,
+ doc_scores: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_retrieved: Optional[bool] = None,
+ do_marginalize: Optional[bool] = None,
+ reduce_loss: Optional[bool] = None,
+ labels: Optional[torch.LongTensor] = None,
+ n_docs: Optional[int] = None,
+ **kwargs, # needs kwargs for generation
+ ) -> RetrievAugLMMarginOutput:
+ r"""
+ do_marginalize (`bool`, *optional*):
+ If `True`, the logits are marginalized over all documents by making use of
+ `torch.nn.functional.log_softmax`.
+ reduce_loss (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `torch.Tensor.sum`
+ operation.
+ kwargs (`Dict[str, any]`, *optional*, defaults to `{}`):
+ Legacy dictionary, which is required so that model can use *generate()* function.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RagRetriever, RagTokenForGeneration
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
+
+ >>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
+ >>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt")
+ >>> input_ids = inputs["input_ids"]
+ >>> labels = targets["input_ids"]
+ >>> outputs = model(input_ids=input_ids, labels=labels)
+
+ >>> # or use retriever separately
+ >>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
+ >>> # 1. Encode
+ >>> question_hidden_states = model.question_encoder(input_ids)[0]
+ >>> # 2. Retrieve
+ >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
+ >>> doc_scores = torch.bmm(
+ ... question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)
+ ... ).squeeze(1)
+ >>> # 3. Forward to generator
+ >>> outputs = model(
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... decoder_input_ids=labels,
+ ... )
+
+ >>> # or directly generate
+ >>> generated = model.generate(
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... )
+ >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
+ ```"""
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ do_marginalize = do_marginalize if do_marginalize is not None else self.config.do_marginalize
+ reduce_loss = reduce_loss if reduce_loss is not None else self.config.reduce_loss
+
+ if labels is not None:
+ if decoder_input_ids is None:
+ decoder_input_ids = labels
+ use_cache = False
+
+ outputs = self.rag(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ doc_scores=doc_scores,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ output_retrieved=output_retrieved,
+ n_docs=n_docs,
+ )
+
+ loss = None
+ logits = outputs.logits
+ if labels is not None:
+ assert decoder_input_ids is not None
+ loss = self.get_nll(
+ outputs.logits,
+ outputs.doc_scores,
+ labels,
+ reduce_loss=reduce_loss,
+ epsilon=self.config.label_smoothing,
+ n_docs=n_docs,
+ )
+
+ if do_marginalize:
+ logits = self.marginalize(logits, outputs.doc_scores, n_docs)
+
+ return RetrievAugLMMarginOutput(
+ loss=loss,
+ logits=logits,
+ doc_scores=outputs.doc_scores,
+ past_key_values=outputs.past_key_values,
+ context_input_ids=outputs.context_input_ids,
+ context_attention_mask=outputs.context_attention_mask,
+ retrieved_doc_embeds=outputs.retrieved_doc_embeds,
+ retrieved_doc_ids=outputs.retrieved_doc_ids,
+ question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state,
+ question_enc_hidden_states=outputs.question_enc_hidden_states,
+ question_enc_attentions=outputs.question_enc_attentions,
+ generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state,
+ generator_enc_hidden_states=outputs.generator_enc_hidden_states,
+ generator_enc_attentions=outputs.generator_enc_attentions,
+ generator_dec_hidden_states=outputs.generator_dec_hidden_states,
+ generator_dec_attentions=outputs.generator_dec_attentions,
+ generator_cross_attentions=outputs.generator_cross_attentions,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.LongTensor] = None,
+ context_input_ids: Optional[torch.LongTensor] = None,
+ context_attention_mask: Optional[torch.LongTensor] = None,
+ doc_scores: Optional[torch.FloatTensor] = None,
+ n_docs: Optional[int] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]] = None,
+ logits_processor: Optional[LogitsProcessorList] = LogitsProcessorList(),
+ stopping_criteria: Optional[StoppingCriteriaList] = StoppingCriteriaList(),
+ **kwargs,
+ ) -> torch.LongTensor:
+ """
+ Implements RAG token decoding.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ The sequence used as a prompt for the generation. If `input_ids` is not passed, then
+ `context_input_ids` has to be provided.
+ 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)
+ context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ n_docs (`int`, *optional*, defaults to `config.n_docs`)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which has the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`, *optional*):
+ If provided, this function constraints the beam search to allowed tokens only at each step. If not
+ provided no constraint is applied. This function takes 2 arguments `inputs_ids` and the batch ID
+ `batch_id`. It has to return a list with the allowed tokens for the next generation step conditioned on
+ the previously generated tokens `inputs_ids` and the batch ID `batch_id`. This argument is useful for
+ constrained generation conditioned on the prefix, as described in [Autoregressive Entity
+ Retrieval](https://arxiv.org/abs/2010.00904).
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and a
+ model's config. If a logit processor is passed that is already created with the arguments or a model's
+ config an error is thrown.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ model's config. If a stopping criteria is passed that is already created with the arguments or a
+ model's config an error is thrown.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model.
+
+ Return:
+ `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated
+ sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter if all batches
+ finished early due to the `eos_token_id`.
+ """
+ # Handle `generation_config` and kwargs that might update it
+ if generation_config is None:
+ generation_config = self.generation_config
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask)
+
+ # set default parameters
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # retrieve docs
+ if self.retriever is not None and context_input_ids is None:
+ question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
+ out = self.retriever(
+ input_ids,
+ question_hidden_states.detach().to(device="cpu", dtype=torch.float32).numpy(),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="pt",
+ )
+ context_input_ids, context_attention_mask, retrieved_doc_embeds = (
+ out["context_input_ids"],
+ out["context_attention_mask"],
+ out["retrieved_doc_embeds"],
+ )
+
+ # set to correct device
+ retrieved_doc_embeds = retrieved_doc_embeds.to(question_hidden_states)
+ context_input_ids = context_input_ids.to(input_ids)
+ context_attention_mask = context_attention_mask.to(input_ids)
+
+ # compute doc_scores
+ doc_scores = torch.bmm(question_hidden_states.unsqueeze(1), retrieved_doc_embeds.transpose(1, 2)).squeeze(
+ 1
+ )
+
+ assert (context_input_ids.shape[0] % n_docs) == 0, (
+ f" The first dimension of `context_input_ids` should be a multiple of `n_docs`={n_docs}, but is"
+ f" {context_input_ids.shape[0]}."
+ )
+
+ # batch_size
+ batch_size = context_input_ids.shape[0] // n_docs
+
+ encoder = self.rag.generator.get_encoder()
+ encoder_outputs = encoder(input_ids=context_input_ids, attention_mask=context_attention_mask, return_dict=True)
+
+ input_ids = torch.full(
+ (batch_size * generation_config.num_beams, 1),
+ generation_config.decoder_start_token_id,
+ dtype=torch.long,
+ device=next(self.parameters()).device,
+ )
+ input_ids_seq_length = input_ids.shape[-1]
+ last_hidden_state = encoder_outputs["last_hidden_state"]
+
+ def extend_enc_output(tensor, num_beams=None):
+ # split into `batch_size`, `num_beams`, `num_docs`
+ tensor = tensor[None, None, :].reshape((batch_size, 1, n_docs) + tensor.shape[1:])
+ # repeat same last hidden states over `num_beams` dimension
+ tensor = tensor.expand((batch_size, num_beams, n_docs) + tensor.shape[3:])
+ # merge `batch_size`, `num_beams`, `num_docs` dims again
+ return tensor.reshape((batch_size * num_beams * n_docs,) + tensor.shape[3:])
+
+ # correctly extend last_hidden_state and attention mask
+ context_attention_mask = extend_enc_output(context_attention_mask, num_beams=generation_config.num_beams)
+ encoder_outputs["last_hidden_state"] = extend_enc_output(
+ last_hidden_state, num_beams=generation_config.num_beams
+ )
+
+ doc_scores = doc_scores.repeat_interleave(generation_config.num_beams, dim=0)
+
+ # define start_len & additional parameters
+ model_kwargs["doc_scores"] = doc_scores
+ model_kwargs["encoder_outputs"] = encoder_outputs
+ model_kwargs["attention_mask"] = context_attention_mask
+ model_kwargs["n_docs"] = n_docs
+
+ pre_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_seq_length,
+ encoder_input_ids=context_input_ids,
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ prepared_stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_config.num_beams == 1:
+ if generation_config.num_return_sequences > 1:
+ raise ValueError(
+ f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"
+ " greedy search."
+ )
+ return self._sample(
+ input_ids,
+ logits_processor=pre_processor,
+ stopping_criteria=prepared_stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=False,
+ streamer=None,
+ **model_kwargs,
+ )
+ elif generation_config.num_beams > 1:
+ if generation_config.num_return_sequences > generation_config.num_beams:
+ raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")
+ return self._beam_search(
+ input_ids,
+ logits_processor=pre_processor,
+ stopping_criteria=prepared_stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=False,
+ **model_kwargs,
+ )
+ else:
+ raise ValueError(
+ f"`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is {generation_config.num_beams}"
+ )
+
+ def get_input_embeddings(self):
+ return self.rag.generator.get_input_embeddings()
+
+ def get_output_embeddings(self):
+ return self.rag.generator.get_output_embeddings()
+
+ def set_output_embeddings(self, new_embeddings):
+ return self.rag.generator.set_output_embeddings(new_embeddings)
+
+ def shift_tokens_right(self, input_ids, start_token_id=None):
+ """Shift input ids one token to the right, and pad with start_token_id"""
+ if start_token_id is None:
+ start_token_id = self.config.decoder_start_token_id
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = start_token_id
+ return shifted_input_ids
+
+ def get_nll(self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, n_docs=None):
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ # shift tokens left
+ target = torch.cat(
+ [target[:, 1:], target.new(target.shape[0], 1).fill_(self.config.generator.pad_token_id)], 1
+ )
+
+ def _mask_pads(ll, smooth_obj):
+ pad_mask = target.eq(self.config.generator.pad_token_id)
+ if pad_mask.any():
+ ll.masked_fill_(pad_mask, 0.0)
+ smooth_obj.masked_fill_(pad_mask, 0.0)
+ return ll.squeeze(-1), smooth_obj.squeeze(-1)
+
+ rag_logprobs = self.marginalize(seq_logits, doc_scores, n_docs)
+
+ target = target.unsqueeze(-1)
+ assert target.dim() == rag_logprobs.dim()
+
+ ll = rag_logprobs.gather(dim=-1, index=target)
+ smooth_obj = rag_logprobs.sum(dim=-1, keepdim=True) # total sum of all (normalised) logits
+ ll, smooth_obj = _mask_pads(ll, smooth_obj)
+ ll = ll.sum(1) # sum over tokens
+ smooth_obj = smooth_obj.sum(1)
+
+ nll_loss = -ll
+ smooth_loss = -smooth_obj
+
+ if reduce_loss:
+ nll_loss = nll_loss.sum()
+ smooth_loss = smooth_loss.sum()
+
+ eps_i = epsilon / rag_logprobs.size(-1)
+ loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
+ return loss
+
+
+__all__ = ["RagModel", "RagPreTrainedModel", "RagSequenceForGeneration", "RagTokenForGeneration"]
diff --git a/lib/python3.12/site-packages/transformers/models/rag/modeling_tf_rag.py b/lib/python3.12/site-packages/transformers/models/rag/modeling_tf_rag.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c670683c9927e8d55bc34eddd642ea266b0dd2a
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/modeling_tf_rag.py
@@ -0,0 +1,1773 @@
+# coding=utf-8
+# Copyright 2020, The RAG 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.
+
+"""TFRAG model implementation."""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import tensorflow as tf
+
+from ...configuration_utils import PretrainedConfig
+from ...generation import TFLogitsProcessorList
+from ...modeling_tf_utils import (
+ TFCausalLanguageModelingLoss,
+ TFModelInputType,
+ TFPreTrainedModel,
+ keras,
+ shape_list,
+ unpack_inputs,
+)
+from ...utils import ModelOutput, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from .configuration_rag import RagConfig
+from .retrieval_rag import RagRetriever
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "RagConfig"
+
+
+@dataclass
+class TFRetrievAugLMMarginOutput(ModelOutput):
+ """
+ Base class for retriever augmented marginalized models outputs.
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss.
+ logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head. The score is possibly marginalized over all documents for
+ each vocabulary token.
+ 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 precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used
+ (see `past_key_values` input) to speed up sequential decoding.
+ doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+ retrieved_doc_embeds (`tf.Tensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*):
+ Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute
+ the `doc_scores`.
+ retrieved_doc_ids (`tf.Tensor` (int32) of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*):
+ The indexes of the embedded documents retrieved by the retriever.
+ context_input_ids (`tf.Tensor`(int32) of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input ids post-processed from the retrieved documents and the question encoder input_ids by the retriever.
+ context_attention_mask (`tf.Tensor` (int32) of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+ question_encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden states at the output of the last layer of the question encoder pooled output of the
+ model.
+ question_enc_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the question encoder at the output of each layer plus the initial embedding outputs.
+ question_enc_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 of the question encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_enc_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the generator encoder of the model.
+ generator_enc_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs.
+ generator_enc_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 of the generator encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_dec_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs.
+ generator_dec_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 of the generator decoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ """
+
+ loss: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ past_key_values: List[tf.Tensor] | None = None
+ doc_scores: tf.Tensor | None = None
+ retrieved_doc_embeds: tf.Tensor | None = None
+ retrieved_doc_ids: tf.Tensor | None = None
+ context_input_ids: tf.Tensor | None = None
+ context_attention_mask: tf.Tensor | None = None
+ question_encoder_last_hidden_state: tf.Tensor | None = None
+ question_enc_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ question_enc_attentions: Tuple[tf.Tensor, ...] | None = None
+ generator_enc_last_hidden_state: tf.Tensor | None = None
+ generator_enc_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ generator_enc_attentions: Tuple[tf.Tensor, ...] | None = None
+ generator_dec_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ generator_dec_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFRetrievAugLMOutput(ModelOutput):
+ """
+ Args:
+ logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head. The score is possibly marginalized over all documents for
+ each vocabulary token.
+ 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 precomputed hidden-states (key and values in the attention blocks) of the decoder that can be used
+ (see `past_key_values` input) to speed up sequential decoding.
+ doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+ retrieved_doc_embeds (`tf.Tensor` of shape `(batch_size, config.n_docs, hidden_size)`, *optional*, returned when *output_retrieved=True*):
+ Embedded documents retrieved by the retriever. Is used with `question_encoder_last_hidden_state` to compute
+ the `doc_scores`.
+ retrieved_doc_ids (`tf.Tensor` of shape `(batch_size, config.n_docs)`, *optional*, returned when *output_retrieved=True*):
+ The indexes of the embedded documents retrieved by the retriever.
+ context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input ids post-processed from the retrieved documents and the question encoder input_ids by the retriever.
+ context_attention_mask (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+ question_encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden states at the output of the last layer of the question encoder pooled output of the
+ model.
+ question_enc_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the question encoder at the output of each layer plus the initial embedding outputs.
+ question_enc_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 of the question encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_enc_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the generator encoder of the model.
+ generator_enc_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator encoder at the output of each layer plus the initial embedding outputs.
+ generator_enc_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 of the generator encoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ generator_dec_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 and one for the output of each layer) of shape
+ `(batch_size, sequence_length, hidden_size)`.
+
+ Hidden states of the generator decoder at the output of each layer plus the initial embedding outputs.
+ generator_dec_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 of the generator decoder, after the attention softmax, used to compute the weighted
+ average in the self-attention heads.
+ """
+
+ logits: Optional[tf.Tensor] = None
+ past_key_values: List[tf.Tensor] | None = None
+ doc_scores: tf.Tensor | None = None
+ retrieved_doc_embeds: tf.Tensor | None = None
+ retrieved_doc_ids: tf.Tensor | None = None
+ context_input_ids: tf.Tensor | None = None
+ context_attention_mask: tf.Tensor | None = None
+ question_encoder_last_hidden_state: tf.Tensor | None = None
+ question_enc_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ question_enc_attentions: Tuple[tf.Tensor, ...] | None = None
+ generator_enc_last_hidden_state: tf.Tensor | None = None
+ generator_enc_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ generator_enc_attentions: Tuple[tf.Tensor, ...] | None = None
+ generator_dec_hidden_states: Tuple[tf.Tensor, ...] | None = None
+ generator_dec_attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+class TFRagPreTrainedModel(TFPreTrainedModel):
+ r"""
+ RAG models were released with the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP
+ Tasks](https://arxiv.org/abs/2005.11401) by Patrick Lewis, Ethan Perez, Aleksandra Piktus et al.
+
+ RAG is a retriever augmented model and encapsulate three components: a question encoder, a dataset retriever and a
+ generator, the encoder and generator are trainable while the retriever is just an indexed dataset.
+
+ """
+
+ config_class = RagConfig
+ base_model_prefix = "rag"
+ _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+ @classmethod
+ def from_pretrained_question_encoder_generator(
+ cls,
+ question_encoder_pretrained_model_name_or_path: Optional[str] = None,
+ generator_pretrained_model_name_or_path: Optional[str] = None,
+ retriever: RagRetriever = None,
+ *model_args,
+ **kwargs,
+ ) -> TFPreTrainedModel:
+ r"""
+ Instantiates an question encoder and a generator from one or two base classes of the library from pretrained
+ model checkpoints.
+
+ Params:
+ question_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the question encoder. Can be either:
+
+ - A string with the *shortcut name* of a pretrained model to load from cache or download, e.g.,
+ `google-bert/bert-base-uncased`.
+ - A string with the *identifier name* of a pretrained model that was user-uploaded to our S3, e.g.,
+ `dbmdz/bert-base-german-cased`.
+ - 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,
+ `question_encoder_from_pt` should be set to `True`.
+
+ generator_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the generator. Can be either:
+
+ - A string with the *shortcut name* of a pretrained model to load from cache or download, e.g.,
+ `google-t5/t5-small`.
+ - A string with the *identifier name* of a pretrained model that was user-uploaded to our S3, e.g.,
+ `facebook/bart-base`.
+ - 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,
+ `generator_from_pt` should be set to `True`.
+
+ model_args (remaining positional arguments, *optional*):
+ All remaining positional arguments will be passed to the underlying model's `__init__` method.
+ retriever ([`RagRetriever`], *optional*):
+ The retriever to use.
+ 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 question_encoder configuration, use the prefix *question_encoder_* for each
+ configuration parameter.
+ - To update the generator configuration, use the prefix *generator_* 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 RagRetriever, TFRagModel
+
+ >>> # initialize a RAG from two pretrained models.
+ >>> model = TFRagModel.from_pretrained_question_encoder_generator(
+ ... "facebook/dpr-question_encoder-single-nq-base", "google-t5/t5-small"
+ ... )
+ >>> # alternatively, initialize from pytorch pretrained models can also be done
+ >>> model = TFRagModel.from_pretrained_question_encoder_generator(
+ ... "facebook/dpr-question_encoder-single-nq-base",
+ ... "facebook/bart-base",
+ ... generator_from_pt=True,
+ ... question_encoder_from_pt=True,
+ ... )
+
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./rag")
+
+ >>> # load retriever
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # load fine-tuned model with retriever
+ >>> model = TFRagModel.from_pretrained("./rag", retriever=retriever)
+ ```"""
+
+ kwargs_question_encoder = {
+ argument[len("question_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("question_encoder_")
+ }
+
+ kwargs_generator = {
+ argument[len("generator_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("generator_")
+ }
+
+ # remove question_encoder, generator kwargs from kwargs
+ for key in kwargs_question_encoder.keys():
+ del kwargs["question_encoder_" + key]
+ for key in kwargs_generator.keys():
+ del kwargs["generator_" + key]
+
+ # Load and initialize the question_encoder and generator
+ # The distinction between question_encoder and generator at the model level is made
+ # by the value of the flag `is_generator` that we need to set correctly.
+ question_encoder = kwargs_question_encoder.pop("model", None)
+ if question_encoder is None:
+ assert question_encoder_pretrained_model_name_or_path is not None, (
+ "If `model` is not defined as an argument, a `question_encoder_pretrained_model_name_or_path` has to"
+ " be defined"
+ )
+
+ from ..auto.modeling_tf_auto import TFAutoModel
+
+ if "config" not in kwargs_question_encoder:
+ from ..auto.configuration_auto import AutoConfig
+
+ question_encoder_config = AutoConfig.from_pretrained(question_encoder_pretrained_model_name_or_path)
+ kwargs_question_encoder["config"] = question_encoder_config
+
+ question_encoder = TFAutoModel.from_pretrained(
+ question_encoder_pretrained_model_name_or_path,
+ name="question_encoder",
+ load_weight_prefix=cls.load_weight_prefix,
+ *model_args,
+ **kwargs_question_encoder,
+ )
+
+ generator = kwargs_generator.pop("generator", None)
+ if generator is None:
+ assert generator_pretrained_model_name_or_path is not None, (
+ "If `generator_model` is not defined as an argument, a `generator_pretrained_model_name_or_path` has"
+ " to be defined"
+ )
+
+ from ..auto.modeling_tf_auto import TFAutoModelForSeq2SeqLM
+
+ if "config" not in kwargs_generator:
+ from ..auto.configuration_auto import AutoConfig
+
+ generator_config = AutoConfig.from_pretrained(generator_pretrained_model_name_or_path)
+ kwargs_generator["config"] = generator_config
+
+ generator = TFAutoModelForSeq2SeqLM.from_pretrained(
+ generator_pretrained_model_name_or_path,
+ name="generator",
+ load_weight_prefix=cls.load_weight_prefix,
+ **kwargs_generator,
+ )
+
+ # instantiate config with corresponding kwargs
+ config = kwargs.get("config", None)
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+
+ return cls(question_encoder=question_encoder, generator=generator, config=config, retriever=retriever)
+
+
+RAG_START_DOCSTRING = r"""
+
+ RAG is a sequence-to-sequence model which encapsulates two core components: a question encoder and a generator.
+ During a forward pass, we encode the input with the question encoder and pass it to the retriever to extract
+ relevant context documents. The documents are then prepended to the input. Such contextualized inputs is passed to
+ the generator.
+
+ The question encoder can be any *autoencoding* model, preferably [`TFDPRQuestionEncoder`], and the generator can be
+ any *seq2seq* model, preferably [`TFBartForConditionalGeneration`].
+
+ The model can be initialized with a [`RagRetriever`] for end-to-end generation or used in combination with the
+ outputs of a retriever in multiple steps---see examples for more details. The model is compatible any
+ *autoencoding* model as the `question_encoder` and any *seq2seq* model with language model head as the `generator`.
+ It has been tested with [`TFDPRQuestionEncoder`] as the `question_encoder` and [`TFBartForConditionalGeneration`]
+ as the `generator`.
+
+ 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.
+
+ The model is in a developing state as it is now fully supports in eager-mode only, and may not be exported in
+ SavedModel format.
+
+ Args:
+ config ([`RagConfig`]):
+ 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.
+ question_encoder ([`TFPreTrainedModel`]):
+ An encoder model compatible with the faiss index encapsulated by the `retriever`.
+ generator ([`TFPreTrainedModel`]):
+ A seq2seq model used as the generator in the RAG architecture.
+ retriever ([`RagRetriever`]):
+ A retriever class encapsulating a faiss index queried to obtain context documents for current inputs.
+"""
+
+
+RAG_FORWARD_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies
+ which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to
+ obtain the indices.
+ attention_mask (`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)
+ encoder_outputs (`tuple(tuple(tf.Tensor)`, *optional*)
+ Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`,
+ *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs *
+ sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the
+ generator's encoder.
+
+ Used by the ([`TFRagModel`]) model during decoding.
+ decoder_input_ids (`tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Provide for generation tasks. `None` by default, construct as per instructions for the generator model
+ you're using with your RAG instance.
+ 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.
+ past_key_values (`tuple(tuple(tf.Tensor))`):
+ Tuple consists of two elements: `encoder_outputs` of the RAG model (see `encoder_outputs`) and
+ `past_key_values` of the underlying generator. Can be used to speed up decoding. `past_key_values` are used
+ in the ([`RagTokenForGeneration`]) model during decoding.
+ doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`. If the model has is not initialized with a `retriever` `doc_scores`
+ has to be provided to the forward pass. `doc_scores` can be computed via
+ `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information.
+ context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model has is not initialized with a `retriever` ``context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`]. context_attention_mask
+ (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when
+ *output_retrieved=True*): Attention mask post-processed from the retrieved documents and the question
+ encoder `input_ids` by the retriever.
+
+ If the model has is not initialized with a `retriever` `context_attention_mask` has to be provided to the
+ forward pass. `context_attention_mask` are returned by [`~RagRetriever.__call__`].
+ 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.
+ output_retrieved(`bool`, *optional*):
+ Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and
+ `context_attention_mask`. See returned tensors for more detail.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`TFRetrievAugLMOutput`] instead of a plain tuple.
+ n_docs (`int`, *optional*, defaults to `config.n_docs``)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+"""
+
+
+@add_start_docstrings_to_model_forward(RAG_START_DOCSTRING)
+class TFRagModel(TFRagPreTrainedModel):
+ load_weight_prefix = "tf_rag_model_1"
+
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[TFPreTrainedModel] = None,
+ generator: Optional[TFPreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None,
+ load_weight_prefix: Optional[str] = None,
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an question_encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+ else:
+ assert isinstance(config, self.config_class), f"config: {config} has to be of type {self.config_class}"
+ super().__init__(config, **kwargs)
+
+ if question_encoder is None:
+ from ..auto.modeling_tf_auto import TFAutoModel
+
+ question_encoder = TFAutoModel.from_config(config.question_encoder, name="question_encoder")
+
+ if generator is None:
+ from ..auto.modeling_tf_auto import TFAutoModelForSeq2SeqLM
+
+ load_weight_prefix = load_weight_prefix if load_weight_prefix is not None else self.load_weight_prefix
+ generator = TFAutoModelForSeq2SeqLM.from_config(
+ config.generator, name="generator", load_weight_prefix=load_weight_prefix + "/generator"
+ )
+
+ self.retriever = retriever
+ if self.retriever is not None:
+ assert isinstance(retriever, RagRetriever), (
+ f"`self.retriever` is of type {type(self.retriever)}, but should be of type `RagRetriever`"
+ )
+ self.retriever = retriever
+
+ self.question_encoder = question_encoder
+ self.generator = generator
+
+ def set_retriever(self, retriever: RagRetriever):
+ self.retriever = retriever
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFRetrievAugLMOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ encoder_outputs: np.ndarray | tf.Tensor | None = None,
+ decoder_input_ids: np.ndarray | tf.Tensor | None = None,
+ decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
+ past_key_values: Tuple[Tuple[Union[np.ndarray, tf.Tensor]]] | None = None,
+ doc_scores: np.ndarray | tf.Tensor | None = None,
+ context_input_ids: np.ndarray | tf.Tensor | None = None,
+ context_attention_mask: np.ndarray | tf.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ output_retrieved: bool | None = None,
+ n_docs: int | None = None,
+ return_dict: bool | None = None,
+ training: bool = False,
+ **kwargs,
+ ) -> TFRetrievAugLMOutput:
+ r"""
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RagRetriever, TFRagModel
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = TFRagModel.from_pretrained("facebook/rag-token-base", retriever=retriever, from_pt=True)
+
+ >>> input_dict = tokenizer.prepare_seq2seq_batch(
+ ... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
+ ... )
+ >>> input_ids = input_dict["input_ids"]
+ >>> outputs = model(input_ids)
+ ```"""
+ assert "decoder_cached_states" not in kwargs, (
+ "Please use past_key_values to cache intermediate outputs"
+ ) # from modeling_tf_bart.py
+
+ # aliasing to minimize code changing
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # whether retriever has to be used
+ has_to_retrieve = (
+ self.retriever is not None
+ and (context_input_ids is None or context_attention_mask is None or doc_scores is None)
+ and encoder_outputs is None
+ )
+
+ # encoder_outputs are pre-computed during RAG-token generation
+ if encoder_outputs is None:
+ if has_to_retrieve:
+ question_enc_outputs = self.question_encoder(
+ input_ids, attention_mask=attention_mask, return_dict=True, training=training
+ )
+ # see https://github.com/huggingface/transformers/blob/main/src/transformers/models/dpr/modeling_tf_dpr.py#L91
+ question_encoder_last_hidden_state = question_enc_outputs[
+ 0
+ ] # hidden states of question encoder => pooler_output
+
+ retriever_outputs = self.retriever(
+ input_ids,
+ question_encoder_last_hidden_state.numpy(),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="tf",
+ )
+ context_input_ids, context_attention_mask, retrieved_doc_embeds, retrieved_doc_ids = (
+ retriever_outputs["context_input_ids"],
+ retriever_outputs["context_attention_mask"],
+ retriever_outputs["retrieved_doc_embeds"],
+ retriever_outputs["doc_ids"],
+ )
+
+ context_input_ids = tf.cast(context_input_ids, tf.int32)
+ context_attention_mask = tf.cast(context_attention_mask, tf.int32)
+ retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)
+ retrieved_doc_ids = tf.cast(retrieved_doc_ids, tf.int32)
+
+ # compute doc_scores
+ doc_scores = tf.squeeze(
+ tf.matmul(
+ tf.expand_dims(question_encoder_last_hidden_state, axis=1),
+ retrieved_doc_embeds,
+ transpose_b=True,
+ ),
+ axis=1,
+ )
+
+ else:
+ assert context_input_ids is not None, (
+ "Make sure that `context_input_ids` are passed, if no `retriever` is set. Alternatively, you can"
+ " set a retriever using the `set_retriever(...)` function."
+ )
+ assert context_attention_mask is not None, (
+ "Make sure that `context_attention_mask` are passed, if no `retriever` is set. Alternatively, you"
+ " can set a retriever using the `set_retriever(...)` function."
+ )
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed, if no `retriever` is set. Alternatively, you can set a"
+ " retriever using the `set_retriever(...)` function."
+ )
+
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed when passing `encoder_outputs` to the forward function."
+ )
+
+ assert (doc_scores.shape[1] % n_docs) == 0, (
+ f" The first dimension of `context_input_ids` should be a multiple of `n_docs`={n_docs}, but is"
+ f" {context_input_ids.shape[0]}."
+ )
+
+ # Decoder input without context documents
+ if decoder_input_ids is not None:
+ decoder_input_ids = tf.repeat(decoder_input_ids, n_docs, axis=0)
+
+ if decoder_attention_mask is not None:
+ decoder_attention_mask = tf.repeat(decoder_attention_mask, n_docs, axis=0)
+
+ gen_outputs = self.generator(
+ context_input_ids,
+ attention_mask=context_attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ return_dict=True,
+ training=training,
+ )
+
+ if not has_to_retrieve:
+ question_encoder_last_hidden_state = None
+ question_enc_hidden_states = None
+ question_enc_attentions = None
+ retrieved_doc_embeds = None
+ retrieved_doc_ids = None
+ else:
+ question_enc_hidden_states = question_enc_outputs.hidden_states
+ question_enc_attentions = question_enc_outputs.attentions
+
+ if not has_to_retrieve or not output_retrieved:
+ # don't output retrieved docs
+ context_input_ids = (None,)
+ context_attention_mask = None
+ retrieved_doc_embeds = None
+ retrieved_doc_ids = None
+
+ return TFRetrievAugLMOutput(
+ logits=gen_outputs.logits,
+ doc_scores=doc_scores,
+ past_key_values=gen_outputs.past_key_values,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ retrieved_doc_embeds=retrieved_doc_embeds,
+ retrieved_doc_ids=retrieved_doc_ids,
+ question_encoder_last_hidden_state=question_encoder_last_hidden_state,
+ question_enc_hidden_states=question_enc_hidden_states,
+ question_enc_attentions=question_enc_attentions,
+ generator_enc_last_hidden_state=gen_outputs.encoder_last_hidden_state,
+ generator_enc_hidden_states=gen_outputs.encoder_hidden_states,
+ generator_enc_attentions=gen_outputs.encoder_attentions,
+ generator_dec_hidden_states=gen_outputs.decoder_hidden_states,
+ generator_dec_attentions=gen_outputs.decoder_attentions,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ with tf.name_scope(self.generator.name):
+ self.generator.build(None)
+ with tf.name_scope(self.question_encoder.name):
+ self.question_encoder.build(None)
+
+
+@add_start_docstrings_to_model_forward(
+ """
+ A TF RAG-token model implementation. It performs RAG-token specific marginalization in the forward pass.
+ """,
+ RAG_START_DOCSTRING,
+)
+class TFRagTokenForGeneration(TFRagPreTrainedModel, TFCausalLanguageModelingLoss):
+ load_weight_prefix = "tf_rag_token_for_generation_1/rag"
+
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[TFPreTrainedModel] = None,
+ generator: Optional[TFPreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None,
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+
+ super().__init__(config)
+
+ # instantiate model
+ self.rag = TFRagModel(
+ config=config,
+ question_encoder=question_encoder,
+ generator=generator,
+ retriever=retriever,
+ load_weight_prefix=self.load_weight_prefix,
+ name="rag",
+ )
+
+ def set_retriever(self, retriever: RagRetriever):
+ self.rag.retriever = retriever
+
+ # Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_bart.py
+ def prepare_inputs_for_generation(
+ self,
+ decoder_input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ use_cache=None,
+ encoder_outputs=None,
+ doc_scores=None,
+ n_docs=None,
+ **kwargs,
+ ):
+ if past_key_values is not None:
+ # if past is defined use only last decoder_input_ids
+ decoder_input_ids = decoder_input_ids[:, -1:]
+
+ return {
+ "input_ids": None,
+ "encoder_outputs": encoder_outputs,
+ "doc_scores": doc_scores,
+ "context_attention_mask": attention_mask,
+ "decoder_input_ids": decoder_input_ids,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ "do_marginalize": True,
+ "n_docs": n_docs,
+ }
+
+ @property
+ def retriever(self):
+ return self.rag.retriever
+
+ @property
+ def generator(self):
+ return self.rag.generator
+
+ @property
+ def question_encoder(self):
+ return self.rag.question_encoder
+
+ @staticmethod
+ def _gather_beams(nested, beam_indices, batch_axis=0):
+ """
+ RAG-specific `_gather_beams`: gathers the beam slices indexed by beam_indices into new beam array. If the
+ nested tensor has a shape mismatch with the beam indices, then it means it is the cache. In that case, isolates
+ and takes care of the extra dimension for ndocs.
+ """
+
+ def gather_fn(tensor):
+ is_rag_cache = tensor.shape[0] != beam_indices.shape[0]
+ if is_rag_cache:
+ n_docs = tensor.shape[0] // beam_indices.shape[0]
+ batch_size = beam_indices.shape[0]
+ # reshapes into (batch size, num beams, n_docs, ...), the cache format expected by RAG
+ tensor = tf.reshape(tensor, (batch_size, -1, n_docs, *tensor.shape[2:]))
+
+ gathered_tensor = tf.gather(params=tensor, indices=beam_indices, axis=1, batch_dims=1)
+
+ if is_rag_cache:
+ # reshapes back into the shape expected by beam search
+ gathered_tensor = tf.reshape(gathered_tensor, (batch_size * n_docs, -1, *gathered_tensor.shape[3:]))
+
+ return gathered_tensor
+
+ return tf.nest.map_structure(gather_fn, nested)
+
+ def marginalize(self, seq_logits, doc_scores, n_docs=None):
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # RAG-token marginalization
+ seq_logprobs = tf.nn.log_softmax(seq_logits, axis=-1)
+ seq_logprobs = tf.reshape(seq_logprobs, [seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.shape[-1]])
+ doc_logprobs = tf.nn.log_softmax(doc_scores, axis=1)
+ doc_logprobs = tf.expand_dims(doc_logprobs, axis=-1)
+ doc_logprobs = tf.expand_dims(doc_logprobs, axis=-1) # twice
+ log_prob_sum = seq_logprobs + doc_logprobs
+ return tf.reduce_logsumexp(log_prob_sum, axis=1)
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFRetrievAugLMMarginOutput, 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[Union[np.ndarray, tf.Tensor]]] | None = None,
+ doc_scores: np.ndarray | tf.Tensor | None = None,
+ context_input_ids: np.ndarray | tf.Tensor | None = None,
+ context_attention_mask: np.ndarray | tf.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ output_retrieved: bool | None = None,
+ n_docs: int | None = None,
+ do_marginalize: bool | None = None,
+ labels: np.ndarray | tf.Tensor | None = None,
+ reduce_loss: bool | None = None,
+ return_dict: bool | None = None,
+ training: bool = False,
+ **kwargs, # needs kwargs for generation
+ ) -> TFRetrievAugLMMarginOutput:
+ r"""
+ do_marginalize (`bool`, *optional*):
+ If `True`, the logits are marginalized over all documents by making use of
+ `torch.nn.functional.log_softmax`.
+ labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the cross entropy classification loss according to Rag-Token model formulation See
+ https://arxiv.org/pdf/2005.11401.pdf Section 2.1 for details about Rag-Token formulation. Indices should be
+ in `[0, ..., config.vocab_size - 1]`.
+ reduce_loss (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `tf.Tensor.sum`
+ operation.
+ kwargs (`Dict[str, any]`, *optional*, defaults to `{}`):
+ Legacy dictionary, which is required so that model can use *generate()* function.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> import tensorflow as tf
+ >>> from transformers import AutoTokenizer, RagRetriever, TFRagTokenForGeneration
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever, from_pt=True)
+
+ >>> input_dict = tokenizer.prepare_seq2seq_batch(
+ ... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
+ ... )
+ >>> outputs = model(input_dict, output_retrieved=True)
+
+ >>> # or use retriever separately
+ >>> # 1. Encode
+ >>> input_ids = input_dict["input_ids"]
+ >>> question_hidden_states = model.question_encoder(input_ids)[0]
+ >>> # 2. Retrieve
+ >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
+ >>> doc_scores = tf.squeeze(
+ ... tf.matmul(
+ ... tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
+ ... ),
+ ... axis=1,
+ ... )
+ >>> # 3. Forward to generator
+ >>> outputs = model(
+ ... inputs=None,
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... decoder_input_ids=input_dict["labels"],
+ ... )
+
+ >>> # or directly generate
+ >>> generated = model.generate(
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... )
+ >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
+ ```"""
+
+ assert "decoder_cached_states" not in kwargs, (
+ "Please use past_key_values to cache intermediate outputs"
+ ) # from modeling_tf_bart.py
+
+ do_marginalize = do_marginalize if do_marginalize else self.config.do_marginalize
+ reduce_loss = reduce_loss if reduce_loss else self.config.reduce_loss
+
+ if labels is not None:
+ if decoder_input_ids is None:
+ decoder_input_ids = labels
+ use_cache = False
+
+ outputs = self.rag(
+ input_ids,
+ attention_mask=attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ doc_scores=doc_scores,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ output_retrieved=output_retrieved,
+ n_docs=n_docs,
+ training=training,
+ )
+
+ loss = None
+ logits = outputs.logits
+ if labels is not None:
+ assert decoder_input_ids is not None
+ loss = self.get_nll(
+ outputs.logits,
+ outputs.doc_scores,
+ labels,
+ reduce_loss=reduce_loss,
+ epsilon=self.config.label_smoothing,
+ n_docs=n_docs,
+ )
+
+ if do_marginalize:
+ logits = self.marginalize(logits, outputs.doc_scores, n_docs)
+
+ return TFRetrievAugLMMarginOutput(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ doc_scores=outputs.doc_scores,
+ context_input_ids=outputs.context_input_ids,
+ context_attention_mask=outputs.context_attention_mask,
+ retrieved_doc_embeds=outputs.retrieved_doc_embeds,
+ retrieved_doc_ids=outputs.retrieved_doc_ids,
+ question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state,
+ question_enc_hidden_states=outputs.question_enc_hidden_states,
+ question_enc_attentions=outputs.question_enc_attentions,
+ generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state,
+ generator_enc_hidden_states=outputs.generator_enc_hidden_states,
+ generator_enc_attentions=outputs.generator_enc_attentions,
+ generator_dec_hidden_states=outputs.generator_dec_hidden_states,
+ generator_dec_attentions=outputs.generator_dec_attentions,
+ )
+
+ def generate(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: tf.Tensor | None = None,
+ context_input_ids=None,
+ context_attention_mask=None,
+ doc_scores=None,
+ n_docs=None,
+ generation_config=None,
+ logits_processor=TFLogitsProcessorList(),
+ **kwargs,
+ ):
+ """
+ Implements TFRAG token decoding.
+
+ Args:
+ input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ The sequence used as a prompt for the generation. If `input_ids` is not passed, then
+ `context_input_ids` has to be provided.
+ attention_mask (`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)
+ context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ context_attention_mask (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`.
+
+ If the model has is not initialized with a `retriever`, `context_input_ids` has to be provided to the
+ forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
+ n_docs (`int`, *optional*, defaults to `config.n_docs`)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`TFLogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and a
+ model's config. If a logit processor is passed that is already created with the arguments or a model's
+ config an error is thrown.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model.
+
+ Return:
+ `tf.Tensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated sequences. The
+ second dimension (sequence_length) is either equal to `max_length` or shorter if all batches finished early
+ due to the `eos_token_id`.
+ """
+ # Handle `generation_config` and kwargs that might update it
+ if generation_config is None:
+ generation_config = self.generation_config
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+
+ # set default parameters
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+
+ # retrieve docs
+ if self.retriever is not None and context_input_ids is None:
+ question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
+ out = self.retriever(
+ input_ids,
+ question_hidden_states.numpy().astype(np.float32),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="tf",
+ )
+ context_input_ids, context_attention_mask, retrieved_doc_embeds = (
+ out["context_input_ids"],
+ out["context_attention_mask"],
+ out["retrieved_doc_embeds"],
+ )
+
+ context_input_ids = tf.cast(context_input_ids, tf.int32)
+ context_attention_mask = tf.cast(context_attention_mask, tf.int32)
+ retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)
+
+ # compute doc_scores
+ doc_scores = tf.matmul(
+ tf.expand_dims(question_hidden_states, axis=1), retrieved_doc_embeds, transpose_b=True
+ )
+ doc_scores = tf.squeeze(doc_scores, axis=1)
+
+ assert (context_input_ids.shape[0] % n_docs) == 0, (
+ f" The first dimension of `context_input_ids` should be a multiple of `n_docs`={n_docs}, but is"
+ f" {context_input_ids.shape[0]}."
+ )
+
+ batch_size = context_input_ids.shape[0] // n_docs
+
+ encoder = self.rag.generator.get_encoder()
+ encoder_outputs = encoder(
+ input_ids=context_input_ids,
+ attention_mask=context_attention_mask,
+ output_attentions=generation_config.output_attentions,
+ output_hidden_states=generation_config.output_hidden_states,
+ return_dict=True,
+ )
+
+ decoder_input_ids = tf.fill(
+ (batch_size * generation_config.num_beams, 1),
+ tf.cast(generation_config.decoder_start_token_id, tf.int32),
+ )
+ last_hidden_state = encoder_outputs["last_hidden_state"]
+
+ def extend_enc_output(tensor, num_beams=None):
+ """
+ Broadcast tensor with `num_beams` replica, with correct order Input: tensor of shape (batch_size*n_docs ,
+ d) Output: tensor of shape (batch_size*num_beams*n_docs , d)
+ """
+
+ # expand batch_size & num_beam dimensions
+ d_shape_list = tensor.shape[1:]
+
+ # split n_docs dimensions
+ new_shape = (batch_size, 1, n_docs) + d_shape_list
+ tensor = tf.reshape(tensor, new_shape)
+
+ # repeat same last hidden states over `num_beams` dimension
+ new_shape = (batch_size, num_beams, n_docs) + d_shape_list
+ tensor = tf.broadcast_to(tensor, new_shape)
+
+ # merge `batch_size`, `num_beams`, `num_docs` dims again
+ new_shape = (batch_size * num_beams * n_docs,) + d_shape_list
+ return tf.reshape(tensor, new_shape)
+
+ # correctly extend last_hidden_state and attention mask
+ context_attention_mask = extend_enc_output(context_attention_mask, num_beams=generation_config.num_beams)
+ encoder_outputs["last_hidden_state"] = extend_enc_output(
+ last_hidden_state, num_beams=generation_config.num_beams
+ )
+
+ doc_scores = tf.repeat(doc_scores, generation_config.num_beams, axis=0)
+
+ # define start_len & additional parameters
+ model_kwargs["doc_scores"] = doc_scores
+ model_kwargs["encoder_outputs"] = encoder_outputs
+ model_kwargs["attention_mask"] = context_attention_mask
+ model_kwargs["n_docs"] = n_docs
+
+ pre_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=tf.shape(decoder_input_ids)[-1],
+ logits_processor=logits_processor,
+ )
+
+ if generation_config.num_beams == 1:
+ return self.greedy_search(
+ input_ids=decoder_input_ids,
+ max_length=generation_config.max_length,
+ pad_token_id=generation_config.pad_token_id,
+ eos_token_id=generation_config.eos_token_id,
+ logits_processor=pre_processor,
+ output_attentions=generation_config.output_attentions,
+ output_hidden_states=generation_config.output_hidden_states,
+ output_scores=generation_config.output_scores,
+ return_dict_in_generate=generation_config.return_dict_in_generate,
+ **model_kwargs,
+ )
+ elif generation_config.num_beams > 1:
+ if generation_config.num_beams < generation_config.num_return_sequences:
+ raise ValueError(
+ "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >="
+ f" num_return_sequences, got {generation_config.num_beams} and"
+ f" {generation_config.num_return_sequences} (respectivelly)"
+ )
+
+ def unflatten_beam_dim(tensor):
+ """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
+ shape = shape_list(tensor)
+ return tf.reshape(tensor, [-1, generation_config.num_beams] + shape[1:])
+
+ decoder_input_ids = unflatten_beam_dim(decoder_input_ids)
+ model_kwargs["attention_mask"] = unflatten_beam_dim(model_kwargs["attention_mask"])
+ model_kwargs["encoder_outputs"]["last_hidden_state"] = unflatten_beam_dim(
+ model_kwargs["encoder_outputs"]["last_hidden_state"]
+ )
+
+ return self.beam_search(
+ input_ids=decoder_input_ids,
+ max_length=generation_config.max_length,
+ pad_token_id=generation_config.pad_token_id,
+ eos_token_id=generation_config.eos_token_id,
+ logits_processor=pre_processor,
+ output_attentions=generation_config.output_attentions,
+ output_hidden_states=generation_config.output_hidden_states,
+ output_scores=generation_config.output_scores,
+ return_dict_in_generate=generation_config.return_dict_in_generate,
+ **model_kwargs,
+ )
+ else:
+ raise ValueError(
+ f"`num_beams` has to be an integer strictly superior to 0 (≥ 1), but is {generation_config.num_beams}"
+ )
+
+ def get_input_embeddings(self):
+ return self.rag.generator.get_input_embeddings()
+
+ def get_output_embeddings(self):
+ return self.rag.generator.get_output_embeddings()
+
+ # Adapted from tf_t5's & tf_bart's _shift_right
+ def shift_tokens_right(self, input_ids, start_token_id=None):
+ """Shift input ids one token to the right, and pad with start_token_id"""
+
+ if start_token_id is None:
+ start_token_id = self.generator.config.decoder_start_token_id
+ assert start_token_id is not None, (
+ "self.generator.config.decoder_start_token_id has to be defined. In Rag we commonly use Bart as"
+ " generator, see Bart docs for more information"
+ )
+
+ pad_token_id = self.generator.config.pad_token_id
+ assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
+
+ start_tokens = tf.fill((shape_list(input_ids)[0], 1), tf.cast(start_token_id, input_ids.dtype))
+ 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), tf.cast(pad_token_id, input_ids.dtype)),
+ shifted_input_ids,
+ )
+
+ # "Verify that `labels` has only positive values and -100"
+ assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, shifted_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
+
+ # nll stands for 'negative log likelihood'
+ def get_nll(self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, n_docs=None):
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ # shift tokens left (from original Pytorch's version)
+
+ target = tf.concat(
+ [target[:, 1:], tf.fill([target.shape[0], 1], tf.cast(self.config.generator.pad_token_id, target.dtype))],
+ axis=1,
+ )
+ rag_logprobs = self.marginalize(seq_logits, doc_scores, n_docs)
+ loss = self.hf_compute_loss(target, rag_logprobs, from_logits=True, reduce_loss=reduce_loss)
+
+ return loss
+
+ # Adopted modeling_tf_bart + add smooth_loss to match with pytorch version
+ def hf_compute_loss(self, labels, y_pred, smooth_epsilon=0.0, from_logits=True, reduce_loss=False):
+ """CrossEntropyLoss that ignores pad tokens"""
+ # Matt: As written, this loss is not XLA-compatible, but it's doing some very weird things
+ # and I don't feel comfortable converting it.
+ loss_fn = keras.losses.SparseCategoricalCrossentropy(
+ from_logits=True,
+ reduction=keras.losses.Reduction.SUM,
+ )
+
+ if from_logits is False: # convert to logits
+ eps = 1e-9
+ y_pred = tf.clip_by_value(y_pred, clip_value_min=eps, clip_value_max=1 - eps)
+ y_pred = tf.math.log(y_pred)
+
+ logits = y_pred
+ melted_labels = tf.reshape(labels, (-1,))
+ active_loss = tf.not_equal(melted_labels, self.config.generator.pad_token_id)
+
+ reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, logits.shape[2])), active_loss)
+ labels = tf.boolean_mask(melted_labels, active_loss)
+ nll_loss = loss_fn(labels, reduced_logits)
+
+ smooth_loss = -tf.reduce_sum(reduced_logits, axis=-1)
+ smooth_loss = tf.reduce_sum(smooth_loss) # sum and squeeze like torch
+ eps_i = smooth_epsilon / reduced_logits.shape[-1]
+
+ loss = (1.0 - smooth_epsilon) * nll_loss + eps_i * smooth_loss
+
+ return loss
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "rag", None) is not None:
+ with tf.name_scope(self.rag.name):
+ self.rag.build(None)
+
+
+@add_start_docstrings_to_model_forward(
+ """
+ A TF RAG-sequence model implementation. It performs RAG-sequence specific marginalization in the forward pass.
+ """,
+ RAG_START_DOCSTRING,
+)
+class TFRagSequenceForGeneration(TFRagPreTrainedModel, TFCausalLanguageModelingLoss):
+ load_weight_prefix = "tf_rag_sequence_for_generation_1/rag"
+
+ def __init__(
+ self,
+ config: Optional[PretrainedConfig] = None,
+ question_encoder: Optional[TFPreTrainedModel] = None,
+ generator: Optional[TFPreTrainedModel] = None,
+ retriever: Optional[RagRetriever] = None,
+ **kwargs,
+ ):
+ assert config is not None or (question_encoder is not None and generator is not None), (
+ "Either a configuration or an encoder and a generator has to be provided."
+ )
+
+ if config is None:
+ config = RagConfig.from_question_encoder_generator_configs(
+ question_encoder.config, generator.config, **kwargs
+ )
+
+ super().__init__(config)
+
+ # instantiate model
+ self.rag = TFRagModel(
+ config=config,
+ question_encoder=question_encoder,
+ generator=generator,
+ retriever=retriever,
+ load_weight_prefix=self.load_weight_prefix,
+ name="rag",
+ )
+
+ def set_retriever(self, retriever: RagRetriever):
+ self.rag.retriever = retriever
+
+ @property
+ def retriever(self):
+ return self.rag.retriever
+
+ @property
+ def generator(self):
+ return self.rag.generator
+
+ @property
+ def question_encoder(self):
+ return self.rag.question_encoder
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(RAG_FORWARD_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFRetrievAugLMMarginOutput, 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: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
+ doc_scores: np.ndarray | tf.Tensor | None = None,
+ context_input_ids: np.ndarray | tf.Tensor | None = None,
+ context_attention_mask: np.ndarray | tf.Tensor | None = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_retrieved: Optional[bool] = None,
+ n_docs: Optional[int] = None,
+ exclude_bos_score: Optional[bool] = None,
+ labels: np.ndarray | tf.Tensor | None = None,
+ reduce_loss: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ training: bool = False,
+ **kwargs, # needs kwargs for generation
+ ) -> Union[Tuple[tf.Tensor], TFRetrievAugLMMarginOutput]:
+ r"""
+ exclude_bos_score (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the score of the BOS token is disregarded when computing
+ the loss.
+ labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the cross entropy classification loss according to Rag-Sequence model formulation See
+ https://arxiv.org/pdf/2005.11401.pdf Section 2.1 for details about Rag-Sequence formulation. Indices should
+ be in `[0, ..., config.vocab_size - 1]`.
+ reduce_loss (`bool`, *optional*):
+ Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `tf.Tensor.sum`
+ operation.
+ kwargs (`Dict[str, any]`, *optional*, defaults to `{}`):
+ Legacy dictionary, which is required so that model can use *generate()* function.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RagRetriever, TFRagSequenceForGeneration
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
+ ... )
+ >>> # initialize with RagRetriever to do everything in one forward call
+ >>> model = TFRagSequenceForGeneration.from_pretrained(
+ ... "facebook/rag-sequence-nq", retriever=retriever, from_pt=True
+ ... )
+
+ >>> input_dict = tokenizer.prepare_seq2seq_batch(
+ ... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
+ ... )
+ >>> outputs = model(input_dict, output_retrieved=True)
+
+ >>> # or use retriever separately
+ >>> # 1. Encode
+ >>> input_ids = input_dict["input_ids"]
+ >>> question_hidden_states = model.question_encoder(input_ids)[0]
+ >>> # 2. Retrieve
+ >>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
+ >>> doc_scores = tf.squeeze(
+ ... tf.matmul(
+ ... tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
+ ... ),
+ ... axis=1,
+ ... )
+ >>> # 3. Forward to generator
+ >>> outputs = model(
+ ... inputs=None,
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... decoder_input_ids=input_dict["labels"],
+ ... )
+
+ >>> # or directly generate
+ >>> generated = model.generate(
+ ... context_input_ids=docs_dict["context_input_ids"],
+ ... context_attention_mask=docs_dict["context_attention_mask"],
+ ... doc_scores=doc_scores,
+ ... )
+ >>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
+ ```"""
+
+ assert "decoder_cached_states" not in kwargs, (
+ "Please use past_key_values to cache intermediate outputs"
+ ) # from modeling_tf_bart.py
+
+ exclude_bos_score = exclude_bos_score if exclude_bos_score else self.config.exclude_bos_score
+ reduce_loss = reduce_loss if reduce_loss else self.config.reduce_loss
+
+ if labels is not None:
+ if decoder_input_ids is None:
+ decoder_input_ids = labels
+ use_cache = False
+
+ outputs = self.rag(
+ input_ids,
+ attention_mask=attention_mask,
+ encoder_outputs=encoder_outputs,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ context_input_ids=context_input_ids,
+ context_attention_mask=context_attention_mask,
+ doc_scores=doc_scores,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ output_retrieved=output_retrieved,
+ n_docs=n_docs,
+ training=training,
+ )
+
+ loss = None
+ if labels is not None:
+ loss = self.get_nll(
+ outputs.logits,
+ outputs.doc_scores,
+ labels,
+ reduce_loss=reduce_loss,
+ epsilon=self.config.label_smoothing,
+ n_docs=n_docs,
+ )
+
+ return TFRetrievAugLMMarginOutput(
+ loss=loss,
+ logits=outputs.logits,
+ doc_scores=outputs.doc_scores,
+ past_key_values=outputs.past_key_values,
+ context_input_ids=outputs.context_input_ids,
+ context_attention_mask=outputs.context_attention_mask,
+ retrieved_doc_embeds=outputs.retrieved_doc_embeds,
+ retrieved_doc_ids=outputs.retrieved_doc_ids,
+ question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state,
+ question_enc_hidden_states=outputs.question_enc_hidden_states,
+ question_enc_attentions=outputs.question_enc_attentions,
+ generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state,
+ generator_enc_hidden_states=outputs.generator_enc_hidden_states,
+ generator_enc_attentions=outputs.generator_enc_attentions,
+ generator_dec_hidden_states=outputs.generator_dec_hidden_states,
+ generator_dec_attentions=outputs.generator_dec_attentions,
+ )
+
+ def get_nll(
+ self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, exclude_bos_score=False, n_docs=None
+ ):
+ # shift tokens left
+ target = tf.concat(
+ [target[:, 1:], tf.fill([target.shape[0], 1], tf.cast(self.config.generator.pad_token_id, target.dtype))],
+ axis=1,
+ )
+
+ # bos_token_id is None for T5
+ bos_token_id = self.config.bos_token_id or self.config.generator.bos_token_id
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ equal_bos_token_id_all = tf.reduce_all(tf.equal(target[:, 0], bos_token_id))
+ use_bos = bos_token_id is not None and equal_bos_token_id_all
+
+ def _mask_pads(ll, smooth_obj):
+ pad_mask = tf.equal(target, tf.cast(self.config.generator.pad_token_id, target.dtype))
+ if tf.reduce_any(pad_mask):
+ ll = tf.where(pad_mask, 0.0, ll)
+ smooth_obj = tf.where(pad_mask, 0.0, smooth_obj)
+ return tf.squeeze(ll, axis=-1), tf.squeeze(smooth_obj, axis=-1)
+
+ # seq_logits.shape = (batch*n_docs, tgt_len , vocabs)
+ seq_logprobs = tf.nn.log_softmax(seq_logits, axis=-1)
+ seq_logprobs = tf.reshape(
+ seq_logprobs, (seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.shape[-1])
+ ) # (batch_size, n_docs, tgt_len, vocabs)
+ doc_logprobs = tf.nn.log_softmax(doc_scores, axis=1)
+ doc_logprobs = tf.expand_dims(doc_logprobs, axis=-1)
+ doc_logprobs = tf.expand_dims(doc_logprobs, axis=-1) # done twice to get 4-D
+
+ # RAG-sequence marginalization
+ first_token_scores = seq_logprobs[:, :, :1, :]
+ second_token_scores = seq_logprobs[:, :, 1:2, :]
+ remainder = seq_logprobs[:, :, 2:, :]
+ rag_logprobs = tf.concat([first_token_scores, second_token_scores + doc_logprobs, remainder], axis=2)
+
+ # calculate loss
+ target = tf.expand_dims(target, axis=1) # n_docs dimension
+ target = tf.expand_dims(target, axis=-1) # logits dimension
+ target = tf.repeat(target, n_docs, axis=1)
+ assert len(target.shape) == len(rag_logprobs.shape)
+
+ # last-axis gathering only - use 2D-reshape-trick for Torch's style nD gathering
+ def torch_gather(param, id_tensor):
+ # 2d-gather torch equivalent: https://stackoverflow.com/questions/52129909/tensorflow-equivalent-of-torch-gather
+ def gather2d(target, id_tensor):
+ idx = tf.stack([tf.range(tf.shape(id_tensor)[0], dtype=id_tensor.dtype), id_tensor[:, 0]], axis=-1)
+ result = tf.gather_nd(target, idx)
+ return tf.expand_dims(result, axis=-1)
+
+ target = tf.reshape(param, (-1, param.shape[-1])) # reshape 2D
+ target_shape = id_tensor.shape
+
+ id_tensor = tf.reshape(id_tensor, (-1, 1)) # also 2D-index
+ result = gather2d(target, id_tensor)
+ return tf.reshape(result, target_shape)
+
+ ll = torch_gather(rag_logprobs, id_tensor=target)
+ smooth_obj = tf.reduce_sum(rag_logprobs, axis=-1, keepdims=True) # total sum of all (normalised) logits
+
+ ll, smooth_obj = _mask_pads(ll, smooth_obj)
+
+ # sum over tokens, exclude bos while scoring
+ if exclude_bos_score and use_bos:
+ ll = tf.reduce_sum(ll[:, :, 1:], axis=2)
+ else:
+ ll = tf.reduce_sum(ll, axis=2)
+
+ smooth_obj = tf.reduce_sum(smooth_obj, axis=2)
+ ll = tf.math.reduce_logsumexp(ll, axis=1) # logsumexp over docs
+ smooth_obj = tf.math.reduce_logsumexp(smooth_obj, axis=1)
+
+ nll_loss = -ll
+ smooth_loss = -smooth_obj
+
+ if reduce_loss:
+ nll_loss = tf.reduce_sum(nll_loss)
+ smooth_loss = tf.reduce_sum(smooth_loss)
+
+ eps_i = epsilon / rag_logprobs.shape[-1]
+ loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
+ return loss
+
+ def generate(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: tf.Tensor | None = None,
+ context_input_ids=None,
+ context_attention_mask=None,
+ doc_scores=None,
+ do_deduplication=None, # defaults to True
+ num_return_sequences=None, # defaults to 1
+ num_beams=None, # defaults to 1
+ n_docs=None,
+ **model_kwargs,
+ ):
+ """
+ Implements RAG sequence "thorough" decoding. Read the [`~generation.GenerationMixin.generate`]` documentation
+ for more information on how to set other generate input parameters
+
+ Args:
+ input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ The sequence used as a prompt for the generation. If `input_ids` is not passed, then
+ `context_input_ids` has to be provided.
+ attention_mask (`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)
+ context_input_ids (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Input IDs post-processed from the retrieved documents and the question encoder input_ids by the
+ retriever.
+ context_attention_mask (`tf.Tensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
+ Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
+ retriever. If the model has is not initialized with a `retriever` or `input_ids` is not given,
+ `context_input_ids` and `context_attention_mask` have to be provided to the forward pass. They are
+ returned by [`~RagRetriever.__call__`].
+ doc_scores (`tf.Tensor` of shape `(batch_size, config.n_docs)`):
+ Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
+ `question_encoder_last_hidden_state`. If the model has is not initialized with a `retriever` or
+ `input_ids` is not given, `doc_scores` has to be provided to the forward pass. `doc_scores` are
+ returned by [`~RagRetriever.__call__`].
+ do_deduplication (`bool`, *optional*):
+ Whether or not to deduplicate the generations from different context documents for a given input. Has
+ to be set to `False` if used while training with distributed backend.
+ num_return_sequences(`int`, *optional*, defaults to 1):
+ The number of independently computed returned sequences for each element in the batch. Note that this
+ is not the value we pass to the `generator`'s `[`~generation.GenerationMixin.generate`]` function,
+ where we set `num_return_sequences` to `num_beams`.
+ num_beams (`int`, *optional*, defaults to 1):
+ Number of beams for beam search. 1 means no beam search.
+ n_docs (`int`, *optional*, defaults to `config.n_docs`)
+ Number of documents to retrieve and/or number of documents for which to generate an answer.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional kwargs will be passed to [`~generation.GenerationMixin.generate`]
+
+ Return:
+ `tf.Tensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated sequences. The
+ second dimension (sequence length) is either equal to `max_length` or shorter if all batches finished early
+ due to the `eos_token_id`.
+ """
+
+ n_docs = n_docs if n_docs is not None else self.config.n_docs
+ do_deduplication = do_deduplication if do_deduplication is not None else self.config.do_deduplication
+ num_doc_return_sequences = (
+ num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
+ )
+ num_beams = num_beams if num_beams is not None else self.config.num_beams
+
+ assert input_ids is not None or context_input_ids is not None, (
+ " At least one of input_ids or context_input_ids must be given"
+ )
+
+ if self.retriever is not None and context_input_ids is None:
+ question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
+ context_input_ids = self.retriever(
+ input_ids,
+ question_hidden_states.numpy(),
+ prefix=self.generator.config.prefix,
+ n_docs=n_docs,
+ return_tensors="tf",
+ )["context_input_ids"]
+
+ hypos = []
+ model_kwargs["num_beams"] = num_beams
+ model_kwargs["num_return_sequences"] = num_beams # put here so that not confused with num_doc_return_sequences
+ model_kwargs["attention_mask"] = None
+
+ batch_size = input_ids.shape[0] if input_ids is not None else context_input_ids.shape[0] // n_docs
+
+ for index in range(batch_size):
+ # first, generate beams from documents:
+ generator_input_ids = context_input_ids[index * n_docs : (index + 1) * n_docs] # (n_docs, max_len)
+
+ output_sequences = self.generator.generate(
+ generator_input_ids,
+ **model_kwargs,
+ ) # n_docs * n_beam, tgt_len
+ if do_deduplication:
+ # do_deduplication -- for TF, work on Eager mode only!
+ output_sequences = tf.stack(list({str(k.numpy().tolist()): k for k in output_sequences}.values()))
+
+ num_candidates = output_sequences.shape[
+ 0
+ ] # after deduplication, this number can be less than n_docs*n_beam
+
+ # then, run model forwards to get nll scores:
+ if input_ids is not None:
+ new_input_ids = tf.tile(input_ids[index : index + 1], (num_candidates, 1))
+ outputs = self(new_input_ids, labels=output_sequences, exclude_bos_score=True)
+ else: # input_ids is None, need context_input_ids/mask and doc_scores
+ assert context_attention_mask is not None, (
+ "Make sure that `context_attention_mask` are passed, if no `input_ids` is set. Alternatively, you"
+ " can set a retriever using the `set_retriever(...)` function."
+ )
+ assert doc_scores is not None, (
+ "Make sure that `doc_scores` are passed, if no `input_ids` is set. Alternatively, you can set a"
+ " retriever using the `set_retriever(...)` function."
+ )
+
+ individual_input_ids = tf.tile(
+ generator_input_ids, (num_candidates, 1)
+ ) # (num_candidates*n_docs, max_len)
+
+ individual_attention_mask = context_attention_mask[index * n_docs : (index + 1) * n_docs]
+ individual_attention_mask = tf.tile(individual_attention_mask, (num_candidates, 1))
+
+ individual_doc_scores = doc_scores[index : (index + 1), :] # doc_scores.shape = [batch, n_docs]
+ individual_doc_scores = tf.tile(individual_doc_scores, (num_candidates, 1)) # [num_candidates, n_docs]
+
+ outputs = self(
+ input_ids=None,
+ context_input_ids=individual_input_ids,
+ context_attention_mask=individual_attention_mask,
+ doc_scores=individual_doc_scores,
+ labels=output_sequences,
+ exclude_bos_score=True,
+ )
+
+ top_cand_inds = tf.math.top_k((-outputs["loss"]), k=num_doc_return_sequences)[1]
+
+ # add hypothesis
+ hypos.append(tf.gather(output_sequences, top_cand_inds))
+
+ return self._cat_and_pad(hypos, pad_token_id=self.config.generator.pad_token_id)
+
+ @staticmethod
+ def _cat_and_pad(tensors, pad_token_id):
+ # used by generate(): tensors is a (batched) list of (candidates, len); len is varied across batch
+
+ # Initialize padded tensor with shape ( all_candidates , max_candidate_length ),
+ # where all_candidates counted from all inputs
+ new_shape = sum([t.shape[0] for t in tensors]), max([t.shape[1] for t in tensors])
+ output = tf.fill(new_shape, pad_token_id)
+
+ # Normal tensor doesn't support slice assignment, so we need tf.Variable
+ output = tf.Variable(output)
+
+ # Assign, and then convert back to tensor
+ ind = 0
+ for t in tensors:
+ output[ind : ind + t.shape[0], : t.shape[1]].assign(t)
+ ind += t.shape[0]
+
+ output = tf.convert_to_tensor(output)
+ return tf.cast(output, tensors[0][0][0].dtype)
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "rag", None) is not None:
+ with tf.name_scope(self.rag.name):
+ self.rag.build(None)
+
+
+__all__ = ["TFRagModel", "TFRagPreTrainedModel", "TFRagSequenceForGeneration", "TFRagTokenForGeneration"]
diff --git a/lib/python3.12/site-packages/transformers/models/rag/retrieval_rag.py b/lib/python3.12/site-packages/transformers/models/rag/retrieval_rag.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7a592a64417f1fa14b3196d0c299958c82fca32
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/retrieval_rag.py
@@ -0,0 +1,677 @@
+# coding=utf-8
+# Copyright 2020, The RAG 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.
+"""RAG Retriever model implementation."""
+
+import os
+import pickle
+import time
+from typing import Iterable, List, Optional, Tuple
+
+import numpy as np
+
+from ...tokenization_utils import PreTrainedTokenizer
+from ...tokenization_utils_base import BatchEncoding
+from ...utils import cached_file, is_datasets_available, is_faiss_available, logging, requires_backends, strtobool
+from .configuration_rag import RagConfig
+from .tokenization_rag import RagTokenizer
+
+
+if is_datasets_available():
+ from datasets import Dataset, load_dataset, load_from_disk
+
+if is_faiss_available():
+ import faiss
+
+
+logger = logging.get_logger(__name__)
+
+
+LEGACY_INDEX_PATH = "https://storage.googleapis.com/huggingface-nlp/datasets/wiki_dpr/"
+
+
+class Index:
+ """
+ A base class for the Indices encapsulated by the [`RagRetriever`].
+ """
+
+ def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
+ """
+ Returns a list of dictionaries, containing titles and text of the retrieved documents.
+
+ Args:
+ doc_ids (`np.ndarray` of shape `(batch_size, n_docs)`):
+ A tensor of document indices.
+ """
+ raise NotImplementedError
+
+ def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+ """
+ For each query in the batch, retrieves `n_docs` documents.
+
+ Args:
+ question_hidden_states (`np.ndarray` of shape `(batch_size, vector_size)`):
+ An array of query vectors.
+ n_docs (`int`):
+ The number of docs retrieved per query.
+
+ Returns:
+ `np.ndarray` of shape `(batch_size, n_docs)`: A tensor of indices of retrieved documents. `np.ndarray` of
+ shape `(batch_size, vector_size)`: A tensor of vector representations of retrieved documents.
+ """
+ raise NotImplementedError
+
+ def is_initialized(self):
+ """
+ Returns `True` if index is already initialized.
+ """
+ raise NotImplementedError
+
+ def init_index(self):
+ """
+ A function responsible for loading the index into memory. Should be called only once per training run of a RAG
+ model. E.g. if the model is trained on multiple GPUs in a distributed setup, only one of the workers will load
+ the index.
+ """
+ raise NotImplementedError
+
+
+class LegacyIndex(Index):
+ """
+ An index which can be deserialized from the files built using https://github.com/facebookresearch/DPR. We use
+ default faiss index parameters as specified in that repository.
+
+ Args:
+ vector_size (`int`):
+ The dimension of indexed vectors.
+ index_path (`str`):
+ A path to a *directory* containing index files compatible with [`~models.rag.retrieval_rag.LegacyIndex`]
+ """
+
+ INDEX_FILENAME = "hf_bert_base.hnswSQ8_correct_phi_128.c_index"
+ PASSAGE_FILENAME = "psgs_w100.tsv.pkl"
+
+ def __init__(self, vector_size, index_path):
+ self.index_id_to_db_id = []
+ self.index_path = index_path
+ self.passages = self._load_passages()
+ self.vector_size = vector_size
+ self.index = None
+ self._index_initialized = False
+
+ def _resolve_path(self, index_path, filename):
+ is_local = os.path.isdir(index_path)
+ try:
+ # Load from URL or cache if already cached
+ resolved_archive_file = cached_file(index_path, filename)
+ except EnvironmentError:
+ msg = (
+ f"Can't load '{filename}'. Make sure that:\n\n"
+ f"- '{index_path}' is a correct remote path to a directory containing a file named {filename}\n\n"
+ f"- or '{index_path}' is the correct path to a directory containing a file named {filename}.\n\n"
+ )
+ raise EnvironmentError(msg)
+ if is_local:
+ logger.info(f"loading file {resolved_archive_file}")
+ else:
+ logger.info(f"loading file {filename} from cache at {resolved_archive_file}")
+ return resolved_archive_file
+
+ def _load_passages(self):
+ logger.info(f"Loading passages from {self.index_path}")
+ passages_path = self._resolve_path(self.index_path, self.PASSAGE_FILENAME)
+ 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(passages_path, "rb") as passages_file:
+ passages = pickle.load(passages_file)
+ return passages
+
+ def _deserialize_index(self):
+ logger.info(f"Loading index from {self.index_path}")
+ resolved_index_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index.dpr")
+ self.index = faiss.read_index(resolved_index_path)
+ resolved_meta_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index_meta.dpr")
+ 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(resolved_meta_path, "rb") as metadata_file:
+ self.index_id_to_db_id = pickle.load(metadata_file)
+ assert len(self.index_id_to_db_id) == self.index.ntotal, (
+ "Deserialized index_id_to_db_id should match faiss index size"
+ )
+
+ def is_initialized(self):
+ return self._index_initialized
+
+ def init_index(self):
+ index = faiss.IndexHNSWFlat(self.vector_size + 1, 512)
+ index.hnsw.efSearch = 128
+ index.hnsw.efConstruction = 200
+ self.index = index
+ self._deserialize_index()
+ self._index_initialized = True
+
+ def get_doc_dicts(self, doc_ids: np.array):
+ doc_list = []
+ for doc_ids_i in doc_ids:
+ ids = [str(int(doc_id)) for doc_id in doc_ids_i]
+ docs = [self.passages[doc_id] for doc_id in ids]
+ doc_list.append(docs)
+ doc_dicts = []
+ for docs in doc_list:
+ doc_dict = {}
+ doc_dict["title"] = [doc[1] for doc in docs]
+ doc_dict["text"] = [doc[0] for doc in docs]
+ doc_dicts.append(doc_dict)
+ return doc_dicts
+
+ def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+ aux_dim = np.zeros(len(question_hidden_states), dtype="float32").reshape(-1, 1)
+ query_nhsw_vectors = np.hstack((question_hidden_states, aux_dim))
+ _, docs_ids = self.index.search(query_nhsw_vectors, n_docs)
+ vectors = [[self.index.reconstruct(int(doc_id))[:-1] for doc_id in doc_ids] for doc_ids in docs_ids]
+ ids = [[int(self.index_id_to_db_id[doc_id]) for doc_id in doc_ids] for doc_ids in docs_ids]
+ return np.array(ids), np.array(vectors)
+
+
+class HFIndexBase(Index):
+ def __init__(self, vector_size, dataset, index_initialized=False):
+ self.vector_size = vector_size
+ self.dataset = dataset
+ self._index_initialized = index_initialized
+ self._check_dataset_format(with_index=index_initialized)
+ dataset.set_format("numpy", columns=["embeddings"], output_all_columns=True, dtype="float32")
+
+ def _check_dataset_format(self, with_index: bool):
+ if not isinstance(self.dataset, Dataset):
+ raise TypeError(f"Dataset should be a datasets.Dataset object, but got {type(self.dataset)}")
+ if len({"title", "text", "embeddings"} - set(self.dataset.column_names)) > 0:
+ raise ValueError(
+ "Dataset should be a dataset with the following columns: "
+ "title (str), text (str) and embeddings (arrays of dimension vector_size), "
+ f"but got columns {self.dataset.column_names}"
+ )
+ if with_index and "embeddings" not in self.dataset.list_indexes():
+ raise ValueError(
+ "Missing faiss index in the dataset. Make sure you called `dataset.add_faiss_index` to compute it "
+ "or `dataset.load_faiss_index` to load one from the disk."
+ )
+
+ def init_index(self):
+ raise NotImplementedError()
+
+ def is_initialized(self):
+ return self._index_initialized
+
+ def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
+ return [self.dataset[doc_ids[i].tolist()] for i in range(doc_ids.shape[0])]
+
+ def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+ _, ids = self.dataset.search_batch("embeddings", question_hidden_states, n_docs)
+ docs = [self.dataset[[i for i in indices if i >= 0]] for indices in ids]
+ vectors = [doc["embeddings"] for doc in docs]
+ for i in range(len(vectors)):
+ if len(vectors[i]) < n_docs:
+ vectors[i] = np.vstack([vectors[i], np.zeros((n_docs - len(vectors[i]), self.vector_size))])
+ return np.array(ids), np.array(vectors) # shapes (batch_size, n_docs) and (batch_size, n_docs, d)
+
+
+class CanonicalHFIndex(HFIndexBase):
+ """
+ A wrapper around an instance of [`~datasets.Datasets`]. If `index_path` is set to `None`, we load the pre-computed
+ index available with the [`~datasets.arrow_dataset.Dataset`], otherwise, we load the index from the indicated path
+ on disk.
+
+ Args:
+ vector_size (`int`): the dimension of the passages embeddings used by the index
+ dataset_name (`str`, optional, defaults to `wiki_dpr`):
+ A dataset identifier of the indexed dataset on HuggingFace AWS bucket (list all available datasets and ids
+ with `datasets.list_datasets()`).
+ dataset_split (`str`, optional, defaults to `train`)
+ Which split of the `dataset` to load.
+ index_name (`str`, optional, defaults to `train`)
+ The index_name of the index associated with the `dataset`. The index loaded from `index_path` will be saved
+ under this name.
+ index_path (`str`, optional, defaults to `None`)
+ The path to the serialized faiss index on disk.
+ use_dummy_dataset (`bool`, optional, defaults to `False`):
+ If True, use the dummy configuration of the dataset for tests.
+ """
+
+ def __init__(
+ self,
+ vector_size: int,
+ dataset_name: str = "wiki_dpr",
+ dataset_split: str = "train",
+ index_name: Optional[str] = None,
+ index_path: Optional[str] = None,
+ use_dummy_dataset=False,
+ dataset_revision=None,
+ ):
+ if int(index_path is None) + int(index_name is None) != 1:
+ raise ValueError("Please provide `index_name` or `index_path`.")
+ self.dataset_name = dataset_name
+ self.dataset_split = dataset_split
+ self.index_name = index_name
+ self.index_path = index_path
+ self.use_dummy_dataset = use_dummy_dataset
+ self.dataset_revision = dataset_revision
+ logger.info(f"Loading passages from {self.dataset_name}")
+ dataset = load_dataset(
+ self.dataset_name,
+ with_index=False,
+ split=self.dataset_split,
+ dummy=self.use_dummy_dataset,
+ revision=dataset_revision,
+ )
+ super().__init__(vector_size, dataset, index_initialized=False)
+
+ def init_index(self):
+ if self.index_path is not None:
+ logger.info(f"Loading index from {self.index_path}")
+ self.dataset.load_faiss_index("embeddings", file=self.index_path)
+ else:
+ logger.info(f"Loading index from {self.dataset_name} with index name {self.index_name}")
+ self.dataset = load_dataset(
+ self.dataset_name,
+ with_embeddings=True,
+ with_index=True,
+ split=self.dataset_split,
+ index_name=self.index_name,
+ dummy=self.use_dummy_dataset,
+ revision=self.dataset_revision,
+ )
+ self.dataset.set_format("numpy", columns=["embeddings"], output_all_columns=True)
+ self._index_initialized = True
+
+
+class CustomHFIndex(HFIndexBase):
+ """
+ A wrapper around an instance of [`~datasets.Datasets`]. The dataset and the index are both loaded from the
+ indicated paths on disk.
+
+ Args:
+ vector_size (`int`): the dimension of the passages embeddings used by the index
+ dataset_path (`str`):
+ The path to the serialized dataset on disk. The dataset should have 3 columns: title (str), text (str) and
+ embeddings (arrays of dimension vector_size)
+ index_path (`str`)
+ The path to the serialized faiss index on disk.
+ """
+
+ def __init__(self, vector_size: int, dataset, index_path=None):
+ super().__init__(vector_size, dataset, index_initialized=index_path is None)
+ self.index_path = index_path
+
+ @classmethod
+ def load_from_disk(cls, vector_size, dataset_path, index_path):
+ logger.info(f"Loading passages from {dataset_path}")
+ if dataset_path is None or index_path is None:
+ raise ValueError(
+ "Please provide `dataset_path` and `index_path` after calling `dataset.save_to_disk(dataset_path)` "
+ "and `dataset.get_index('embeddings').save(index_path)`."
+ )
+ dataset = load_from_disk(dataset_path)
+ return cls(vector_size=vector_size, dataset=dataset, index_path=index_path)
+
+ def init_index(self):
+ if not self.is_initialized():
+ logger.info(f"Loading index from {self.index_path}")
+ self.dataset.load_faiss_index("embeddings", file=self.index_path)
+ self._index_initialized = True
+
+
+class RagRetriever:
+ """
+ Retriever used to get documents from vector queries. It retrieves the documents embeddings as well as the documents
+ contents, and it formats them to be used with a RagModel.
+
+ Args:
+ config ([`RagConfig`]):
+ The configuration of the RAG model this Retriever is used with. Contains parameters indicating which
+ `Index` to build. You can load your own custom dataset with `config.index_name="custom"` or use a canonical
+ one (default) from the datasets library with `config.index_name="wiki_dpr"` for example.
+ question_encoder_tokenizer ([`PreTrainedTokenizer`]):
+ The tokenizer that was used to tokenize the question. It is used to decode the question and then use the
+ generator_tokenizer.
+ generator_tokenizer ([`PreTrainedTokenizer`]):
+ The tokenizer used for the generator part of the RagModel.
+ index ([`~models.rag.retrieval_rag.Index`], optional, defaults to the one defined by the configuration):
+ If specified, use this index instead of the one built using the configuration
+
+ Examples:
+
+ ```python
+ >>> # To load the default "wiki_dpr" dataset with 21M passages from wikipedia (index name is 'compressed' or 'exact')
+ >>> from transformers import RagRetriever
+
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/dpr-ctx_encoder-single-nq-base", dataset="wiki_dpr", index_name="compressed"
+ ... )
+
+ >>> # To load your own indexed dataset built with the datasets library. More info on how to build the indexed dataset in examples/rag/use_own_knowledge_dataset.py
+ >>> from transformers import RagRetriever
+
+ >>> dataset = (
+ ... ...
+ ... ) # dataset must be a datasets.Datasets object with columns "title", "text" and "embeddings", and it must have a faiss index
+ >>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", indexed_dataset=dataset)
+
+ >>> # To load your own indexed dataset built with the datasets library that was saved on disk. More info in examples/rag/use_own_knowledge_dataset.py
+ >>> from transformers import RagRetriever
+
+ >>> dataset_path = "path/to/my/dataset" # dataset saved via *dataset.save_to_disk(...)*
+ >>> index_path = "path/to/my/index.faiss" # faiss index saved via *dataset.get_index("embeddings").save(...)*
+ >>> retriever = RagRetriever.from_pretrained(
+ ... "facebook/dpr-ctx_encoder-single-nq-base",
+ ... index_name="custom",
+ ... passages_path=dataset_path,
+ ... index_path=index_path,
+ ... )
+
+ >>> # To load the legacy index built originally for Rag's paper
+ >>> from transformers import RagRetriever
+
+ >>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", index_name="legacy")
+ ```"""
+
+ def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, index=None, init_retrieval=True):
+ self._init_retrieval = init_retrieval
+ requires_backends(self, ["datasets", "faiss"])
+ super().__init__()
+ self.index = index or self._build_index(config)
+ self.generator_tokenizer = generator_tokenizer
+ self.question_encoder_tokenizer = question_encoder_tokenizer
+
+ self.n_docs = config.n_docs
+ self.batch_size = config.retrieval_batch_size
+
+ self.config = config
+ if self._init_retrieval:
+ self.init_retrieval()
+
+ self.ctx_encoder_tokenizer = None
+ self.return_tokenized_docs = False
+
+ @staticmethod
+ def _build_index(config):
+ if config.index_name == "legacy":
+ return LegacyIndex(
+ config.retrieval_vector_size,
+ config.index_path or LEGACY_INDEX_PATH,
+ )
+ elif config.index_name == "custom":
+ return CustomHFIndex.load_from_disk(
+ vector_size=config.retrieval_vector_size,
+ dataset_path=config.passages_path,
+ index_path=config.index_path,
+ )
+ else:
+ return CanonicalHFIndex(
+ vector_size=config.retrieval_vector_size,
+ dataset_name=config.dataset,
+ dataset_split=config.dataset_split,
+ index_name=config.index_name,
+ index_path=config.index_path,
+ use_dummy_dataset=config.use_dummy_dataset,
+ dataset_revision=config.dataset_revision,
+ )
+
+ @classmethod
+ def from_pretrained(cls, retriever_name_or_path, indexed_dataset=None, **kwargs):
+ requires_backends(cls, ["datasets", "faiss"])
+ config = kwargs.pop("config", None) or RagConfig.from_pretrained(retriever_name_or_path, **kwargs)
+ rag_tokenizer = RagTokenizer.from_pretrained(retriever_name_or_path, config=config)
+ question_encoder_tokenizer = rag_tokenizer.question_encoder
+ generator_tokenizer = rag_tokenizer.generator
+ if indexed_dataset is not None:
+ config.index_name = "custom"
+ index = CustomHFIndex(config.retrieval_vector_size, indexed_dataset)
+ else:
+ index = cls._build_index(config)
+ return cls(
+ config,
+ question_encoder_tokenizer=question_encoder_tokenizer,
+ generator_tokenizer=generator_tokenizer,
+ index=index,
+ )
+
+ def save_pretrained(self, save_directory):
+ if isinstance(self.index, CustomHFIndex):
+ if self.config.index_path is None:
+ index_path = os.path.join(save_directory, "hf_dataset_index.faiss")
+ self.index.dataset.get_index("embeddings").save(index_path)
+ self.config.index_path = index_path
+ if self.config.passages_path is None:
+ passages_path = os.path.join(save_directory, "hf_dataset")
+ # datasets don't support save_to_disk with indexes right now
+ faiss_index = self.index.dataset._indexes.pop("embeddings")
+ self.index.dataset.save_to_disk(passages_path)
+ self.index.dataset._indexes["embeddings"] = faiss_index
+ self.config.passages_path = passages_path
+ self.config.save_pretrained(save_directory)
+ rag_tokenizer = RagTokenizer(
+ question_encoder=self.question_encoder_tokenizer,
+ generator=self.generator_tokenizer,
+ )
+ rag_tokenizer.save_pretrained(save_directory)
+
+ def init_retrieval(self):
+ """
+ Retriever initialization function. It loads the index into memory.
+ """
+
+ logger.info("initializing retrieval")
+ self.index.init_index()
+
+ def postprocess_docs(self, docs, input_strings, prefix, n_docs, return_tensors=None):
+ r"""
+ Postprocessing retrieved `docs` and combining them with `input_strings`.
+
+ Args:
+ docs (`dict`):
+ Retrieved documents.
+ input_strings (`str`):
+ Input strings decoded by `preprocess_query`.
+ prefix (`str`):
+ Prefix added at the beginning of each input, typically used with T5-based models.
+
+ Return:
+ `tuple(tensors)`: a tuple consisting of two elements: contextualized `input_ids` and a compatible
+ `attention_mask`.
+ """
+
+ def cat_input_and_doc(doc_title, doc_text, input_string, prefix):
+ # TODO(Patrick): if we train more RAG models, I want to put the input first to take advantage of effortless truncation
+ # TODO(piktus): better handling of truncation
+ if doc_title.startswith('"'):
+ doc_title = doc_title[1:]
+ if doc_title.endswith('"'):
+ doc_title = doc_title[:-1]
+ if prefix is None:
+ prefix = ""
+ out = (prefix + doc_title + self.config.title_sep + doc_text + self.config.doc_sep + input_string).replace(
+ " ", " "
+ )
+ return out
+
+ rag_input_strings = [
+ cat_input_and_doc(
+ docs[i]["title"][j],
+ docs[i]["text"][j],
+ input_strings[i],
+ prefix,
+ )
+ for i in range(len(docs))
+ for j in range(n_docs)
+ ]
+
+ contextualized_inputs = self.generator_tokenizer.batch_encode_plus(
+ rag_input_strings,
+ max_length=self.config.max_combined_length,
+ return_tensors=return_tensors,
+ padding="max_length",
+ truncation=True,
+ )
+
+ return contextualized_inputs["input_ids"], contextualized_inputs["attention_mask"]
+
+ def _chunk_tensor(self, t: Iterable, chunk_size: int) -> List[Iterable]:
+ return [t[i : i + chunk_size] for i in range(0, len(t), chunk_size)]
+
+ def _main_retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, np.ndarray]:
+ question_hidden_states_batched = self._chunk_tensor(question_hidden_states, self.batch_size)
+ ids_batched = []
+ vectors_batched = []
+ for question_hidden_states in question_hidden_states_batched:
+ start_time = time.time()
+ ids, vectors = self.index.get_top_docs(question_hidden_states, n_docs)
+ logger.debug(
+ f"index search time: {time.time() - start_time} sec, batch size {question_hidden_states.shape}"
+ )
+ ids_batched.extend(ids)
+ vectors_batched.extend(vectors)
+ return (
+ np.array(ids_batched),
+ np.array(vectors_batched),
+ ) # shapes (batch_size, n_docs) and (batch_size, n_docs, d)
+
+ def retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, List[dict]]:
+ """
+ Retrieves documents for specified `question_hidden_states`.
+
+ Args:
+ question_hidden_states (`np.ndarray` of shape `(batch_size, vector_size)`):
+ A batch of query vectors to retrieve with.
+ n_docs (`int`):
+ The number of docs retrieved per query.
+
+ Return:
+ `Tuple[np.ndarray, np.ndarray, List[dict]]`: A tuple with the following objects:
+
+ - **retrieved_doc_embeds** (`np.ndarray` of shape `(batch_size, n_docs, dim)`) -- The retrieval embeddings
+ of the retrieved docs per query.
+ - **doc_ids** (`np.ndarray` of shape `(batch_size, n_docs)`) -- The ids of the documents in the index
+ - **doc_dicts** (`List[dict]`): The `retrieved_doc_embeds` examples per query.
+ """
+
+ doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
+ return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)
+
+ def set_ctx_encoder_tokenizer(self, ctx_encoder_tokenizer: PreTrainedTokenizer):
+ # used in end2end retriever training
+ self.ctx_encoder_tokenizer = ctx_encoder_tokenizer
+ self.return_tokenized_docs = True
+
+ def __call__(
+ self,
+ question_input_ids: List[List[int]],
+ question_hidden_states: np.ndarray,
+ prefix=None,
+ n_docs=None,
+ return_tensors=None,
+ ) -> BatchEncoding:
+ """
+ Retrieves documents for specified `question_hidden_states`.
+
+ Args:
+ question_input_ids (`List[List[int]]`) batch of input ids
+ question_hidden_states (`np.ndarray` of shape `(batch_size, vector_size)`:
+ A batch of query vectors to retrieve with.
+ prefix (`str`, *optional*):
+ The prefix used by the generator's tokenizer.
+ n_docs (`int`, *optional*):
+ The number of docs retrieved per query.
+ return_tensors (`str` or [`~utils.TensorType`], *optional*, defaults to "pt"):
+ 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.
+
+ Returns: [`BatchEncoding`]: A [`BatchEncoding`] with the following fields:
+
+ - **context_input_ids** -- List of token ids to be fed to a model.
+
+ [What are input IDs?](../glossary#input-ids)
+
+ - **context_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`).
+
+ [What are attention masks?](../glossary#attention-mask)
+
+ - **retrieved_doc_embeds** -- List of embeddings of the retrieved documents
+ - **doc_ids** -- List of ids of the retrieved documents
+ """
+
+ n_docs = n_docs if n_docs is not None else self.n_docs
+ prefix = prefix if prefix is not None else self.config.generator.prefix
+ retrieved_doc_embeds, doc_ids, docs = self.retrieve(question_hidden_states, n_docs)
+
+ input_strings = self.question_encoder_tokenizer.batch_decode(question_input_ids, skip_special_tokens=True)
+ context_input_ids, context_attention_mask = self.postprocess_docs(
+ docs, input_strings, prefix, n_docs, return_tensors=return_tensors
+ )
+
+ if self.return_tokenized_docs:
+ retrieved_doc_text = []
+ retrieved_doc_title = []
+
+ for b_idx in range(len(docs)):
+ for doc_idx in range(n_docs):
+ retrieved_doc_text.append(docs[b_idx]["text"][doc_idx])
+ retrieved_doc_title.append(docs[b_idx]["title"][doc_idx])
+
+ tokenized_docs = self.ctx_encoder_tokenizer(
+ retrieved_doc_title,
+ retrieved_doc_text,
+ truncation=True,
+ padding="longest",
+ return_tensors=return_tensors,
+ )
+
+ return BatchEncoding(
+ {
+ "context_input_ids": context_input_ids,
+ "context_attention_mask": context_attention_mask,
+ "retrieved_doc_embeds": retrieved_doc_embeds,
+ "doc_ids": doc_ids,
+ "tokenized_doc_ids": tokenized_docs["input_ids"],
+ "tokenized_doc_attention_mask": tokenized_docs["attention_mask"],
+ },
+ tensor_type=return_tensors,
+ )
+
+ else:
+ return BatchEncoding(
+ {
+ "context_input_ids": context_input_ids,
+ "context_attention_mask": context_attention_mask,
+ "retrieved_doc_embeds": retrieved_doc_embeds,
+ "doc_ids": doc_ids,
+ },
+ tensor_type=return_tensors,
+ )
+
+
+__all__ = ["RagRetriever"]
diff --git a/lib/python3.12/site-packages/transformers/models/rag/tokenization_rag.py b/lib/python3.12/site-packages/transformers/models/rag/tokenization_rag.py
new file mode 100644
index 0000000000000000000000000000000000000000..428b5f66446c022fb37d34eecd4c360b720c0cc7
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/rag/tokenization_rag.py
@@ -0,0 +1,124 @@
+# coding=utf-8
+# Copyright 2020, The RAG 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 RAG."""
+
+import os
+import warnings
+from typing import List, Optional
+
+from ...tokenization_utils_base import BatchEncoding
+from ...utils import logging
+from .configuration_rag import RagConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class RagTokenizer:
+ def __init__(self, question_encoder, generator):
+ self.question_encoder = question_encoder
+ self.generator = generator
+ self.current_tokenizer = self.question_encoder
+
+ def save_pretrained(self, save_directory):
+ if os.path.isfile(save_directory):
+ raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+ os.makedirs(save_directory, exist_ok=True)
+ question_encoder_path = os.path.join(save_directory, "question_encoder_tokenizer")
+ generator_path = os.path.join(save_directory, "generator_tokenizer")
+ self.question_encoder.save_pretrained(question_encoder_path)
+ self.generator.save_pretrained(generator_path)
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ # dynamically import AutoTokenizer
+ from ..auto.tokenization_auto import AutoTokenizer
+
+ config = kwargs.pop("config", None)
+
+ if config is None:
+ config = RagConfig.from_pretrained(pretrained_model_name_or_path)
+
+ question_encoder = AutoTokenizer.from_pretrained(
+ pretrained_model_name_or_path, config=config.question_encoder, subfolder="question_encoder_tokenizer"
+ )
+ generator = AutoTokenizer.from_pretrained(
+ pretrained_model_name_or_path, config=config.generator, subfolder="generator_tokenizer"
+ )
+ return cls(question_encoder=question_encoder, generator=generator)
+
+ def __call__(self, *args, **kwargs):
+ return self.current_tokenizer(*args, **kwargs)
+
+ def batch_decode(self, *args, **kwargs):
+ return self.generator.batch_decode(*args, **kwargs)
+
+ def decode(self, *args, **kwargs):
+ return self.generator.decode(*args, **kwargs)
+
+ def _switch_to_input_mode(self):
+ self.current_tokenizer = self.question_encoder
+
+ def _switch_to_target_mode(self):
+ self.current_tokenizer = self.generator
+
+ def prepare_seq2seq_batch(
+ self,
+ src_texts: List[str],
+ tgt_texts: Optional[List[str]] = None,
+ max_length: Optional[int] = None,
+ max_target_length: Optional[int] = None,
+ padding: str = "longest",
+ return_tensors: Optional[str] = None,
+ truncation: bool = True,
+ **kwargs,
+ ) -> BatchEncoding:
+ warnings.warn(
+ "`prepare_seq2seq_batch` is deprecated and will be removed in version 5 of 🤗 Transformers. Use the "
+ "regular `__call__` method to prepare your inputs and the tokenizer under the `with_target_tokenizer` "
+ "context manager to prepare your targets. See the documentation of your specific tokenizer for more "
+ "details",
+ FutureWarning,
+ )
+ if max_length is None:
+ max_length = self.current_tokenizer.model_max_length
+ model_inputs = self(
+ src_texts,
+ add_special_tokens=True,
+ return_tensors=return_tensors,
+ max_length=max_length,
+ padding=padding,
+ truncation=truncation,
+ **kwargs,
+ )
+ if tgt_texts is None:
+ return model_inputs
+ # Process tgt_texts
+ if max_target_length is None:
+ max_target_length = self.current_tokenizer.model_max_length
+ labels = self(
+ text_target=tgt_texts,
+ add_special_tokens=True,
+ return_tensors=return_tensors,
+ padding=padding,
+ max_length=max_target_length,
+ truncation=truncation,
+ **kwargs,
+ )
+ model_inputs["labels"] = labels["input_ids"]
+ return model_inputs
+
+
+__all__ = ["RagTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__init__.py b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab9335fc4beff8b178336ccb546bdecb4cd45171
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_recurrent_gemma import *
+ from .modeling_recurrent_gemma import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e5288ec5aa9fef3190d4bc41d7534c0784d82963
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/configuration_recurrent_gemma.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/configuration_recurrent_gemma.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9f08174bc46848105593d1c6378408a27f267144
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/configuration_recurrent_gemma.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/modeling_recurrent_gemma.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/modeling_recurrent_gemma.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e507017635125d5cd795dd2a8817c158a2590cac
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/__pycache__/modeling_recurrent_gemma.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py
new file mode 100644
index 0000000000000000000000000000000000000000..60a034f57d3dcfcf3398a79854363f9400c44db2
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py
@@ -0,0 +1,161 @@
+# 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.
+"""RecurrentGemma model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class RecurrentGemmaConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`RecurrentGemmaModel`]. It is used to instantiate a RecurrentGemma
+ 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 RecurrentGemma-7B.
+
+ e.g. [google/recurrentgemma-2b](https://huggingface.co/google/recurrentgemma-2b)
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ num_hidden_layers (`int`, *optional*, defaults to 26):
+ The number of hidden layers in the model.
+ vocab_size (`int`, *optional*, defaults to 256000):
+ Vocabulary size of the RecurrentGemma model. Defines the number of
+ different tokens that can be represented by the
+ `inputs_ids` passed when calling [`RecurrentGemmaModel`]
+ hidden_size (`int`, *optional*, defaults to 2560):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 7680):
+ Dimension of the MLP representations.
+ num_attention_heads (`int`, *optional*, defaults to 10):
+ The number of heads for the attention block and the number of
+ heads/blocks for the block-diagonal layers used in the RG-LRU gates.
+ This number must divide `hidden_size` and `lru_width`.
+ lru_width (`int` or `None`, *optional*):
+ Dimension of the hidden representations of the RG-LRU. If `None`
+ this will be set to `hidden_size`.
+ Whether to scale the output of the embeddings by `sqrt(hidden_size)`.
+ attention_window_size (`int`, *optional*, defaults to 2048):
+ The size of the attention window used in the attention block.
+ conv1d_width (`int`, *optional*, defaults to 4):
+ The kernel size of conv1d layers used in the recurrent blocks.
+ logits_soft_cap (`float`, *optional*, defaults to 30.0):
+ The value at which the logits should be soft-capped to after the transformer and LM-head computation in the Causal LM architecture.
+ 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 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.
+ hidden_activation (``str` or `function``, *optional*, defaults to `"gelu_pytorch_tanh"`):
+ The hidden activation used in the recurrent block as well as the MLP layer of the decoder layers.
+ partial_rotary_factor (`float`, *optional*, defaults to 0.5):
+ The partial rotary factor used in the initialization of the rotary embeddings.
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ block_types (`List[str]`, *optional*, defaults to `('recurrent', 'recurrent', 'attention')`):
+ List of aleternating blocks that will be repeated to initialize the `temporal_block` layer.
+ attention_dropout (`float`, *optional*, defaults to 0.0): dropout value to use after the attention softmax.
+ num_key_value_heads (`16`, *optional*, defaults to 16): Number of key value heads to use GQA.
+ attention_bias (`bool`, *optional*, defaults to `False`): whether or not the linear q,k,v of the Attention layer should have bias
+ w_init_variance_scale (`float`, *optional*, defaults to 0.01): weight initialization variance.
+ ```python
+ >>> from transformers import RecurrentGemmaModel, RecurrentGemmaConfig
+
+ >>> # Initializing a RecurrentGemma recurrentgemma-2b style configuration
+ >>> configuration = RecurrentGemmaConfig()
+
+ >>> # Initializing a model from the recurrentgemma-2b style configuration
+ >>> model = RecurrentGemmaModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "recurrent_gemma"
+
+ def __init__(
+ self,
+ num_hidden_layers=26,
+ vocab_size=256000,
+ hidden_size=2560,
+ intermediate_size=3 * 2560,
+ num_attention_heads=10,
+ lru_width=None,
+ attention_window_size=2048,
+ conv1d_width=4,
+ logits_soft_cap=30.0,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ pad_token_id=0,
+ eos_token_id=1,
+ bos_token_id=2,
+ hidden_activation="gelu_pytorch_tanh",
+ partial_rotary_factor=0.5,
+ rope_theta=10000.0,
+ block_types=("recurrent", "recurrent", "attention"),
+ attention_dropout=0.0,
+ num_key_value_heads=None,
+ attention_bias=False,
+ w_init_variance_scale=0.01,
+ **kwargs,
+ ):
+ self.num_hidden_layers = num_hidden_layers
+ self.vocab_size = vocab_size
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_attention_heads = num_attention_heads
+ self.lru_width = lru_width if lru_width is not None else hidden_size
+ self.attention_window_size = attention_window_size
+ self.conv1d_width = conv1d_width
+ self.logits_soft_cap = logits_soft_cap
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.partial_rotary_factor = partial_rotary_factor
+ self.block_types = list(block_types)
+ self.hidden_activation = hidden_activation
+ self.head_dim = self.hidden_size // self.num_attention_heads
+ self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+ if self.num_key_value_heads > self.num_attention_heads:
+ raise ValueError("The number of `num_key_value_heads` must be smaller than `num_attention_heads`")
+ self.attention_dropout = attention_dropout
+ self.attention_bias = attention_bias
+ self.w_init_variance_scale = w_init_variance_scale
+ self.final_w_init_variance_scale = 2.0 / self.num_hidden_layers
+ super().__init__(
+ pad_token_id=pad_token_id,
+ bos_token_id=bos_token_id,
+ eos_token_id=eos_token_id,
+ **kwargs,
+ )
+
+ @property
+ def layers_block_type(self):
+ return (self.block_types * 100)[: self.num_hidden_layers]
+
+
+__all__ = ["RecurrentGemmaConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
new file mode 100644
index 0000000000000000000000000000000000000000..450da13493d5ea0b42b26c58cb183b1fd4df0fb5
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
@@ -0,0 +1,903 @@
+# 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 RecurrentGemma model."""
+
+import math
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import AttentionMaskConverter
+from ...modeling_outputs import BaseModelOutputWithNoAttention, CausalLMOutput
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import ALL_LAYERNORM_LAYERS
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.import_utils import is_torchdynamo_compiling
+from .configuration_recurrent_gemma import RecurrentGemmaConfig
+
+
+logger = logging.get_logger(__name__)
+_CONFIG_FOR_DOC = "RecurrentGemmaConfig"
+_MAX_SQRT_GRADIENT = 1000.0
+
+
+# Copied from transformers.models.gemma.modeling_gemma.GemmaRMSNorm with Gemma->RecurrentGemma
+class RecurrentGemmaRMSNorm(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 RecurrentGemma 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+ALL_LAYERNORM_LAYERS.append(RecurrentGemmaRMSNorm)
+
+
+class RecurrentGemmaRotaryEmbedding(nn.Module):
+ def __init__(self, dim, base=10000, device=None):
+ super().__init__()
+ self.dim = dim
+ self.base = base
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+ self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)
+
+ @torch.no_grad()
+ def forward(self, x, position_ids, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ self.inv_freq.to(x.device)
+ 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)
+
+
+class RecurrentGemmaSdpaAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: RecurrentGemmaConfig):
+ super().__init__()
+ self.config = config
+ self.attention_dropout = config.attention_dropout
+ self.hidden_size = config.hidden_size
+ self.num_attention_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_attention_heads // self.num_key_value_heads
+ self.partial_rotary_factor = config.partial_rotary_factor
+
+ self.q_proj = nn.Linear(self.hidden_size, self.num_attention_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_attention_heads * self.head_dim, self.hidden_size, bias=True)
+ self.rotary_emb = RecurrentGemmaRotaryEmbedding(
+ int(self.partial_rotary_factor * self.head_dim),
+ base=config.rope_theta,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ 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)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, self.num_attention_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)
+
+ # Partial rotary embedding
+ query_rot, query_pass = torch.chunk(query_states, int(1 / self.partial_rotary_factor), dim=-1)
+ key_rot, key_pass = torch.chunk(key_states, int(1 / self.partial_rotary_factor), dim=-1)
+ query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
+
+ if use_cache and hasattr(self, "key_states"):
+ cache_kwargs = {"cache_position": cache_position}
+ key_states, value_states = self._update_cache(key_states, value_states, **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]]
+
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_states.contiguous(),
+ key_states.contiguous(),
+ value_states.contiguous(),
+ attn_mask=causal_mask, # pretty much a must for sliding window backend!
+ dropout_p=self.attention_dropout if self.training else 0.0,
+ scale=self.head_dim**-0.5,
+ )
+
+ 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
+
+ def _setup_cache(self, batch_size, device, dtype=None):
+ if dtype is None and self.config.torch_dtype is not None:
+ dtype = self.config.torch_dtype
+ dtype = dtype if dtype is not None else torch.float32
+ cache_shape = (batch_size, self.num_key_value_heads, self.config.attention_window_size, self.head_dim)
+ self.value_states = torch.zeros(cache_shape, dtype=dtype, device=device)
+ self.key_states = torch.zeros(cache_shape, dtype=dtype, device=device)
+
+ @torch.no_grad()
+ def _update_cache(self, key_states, value_states, **cache_kwargs):
+ """
+ torch.compile compatible sliding window.
+ Computes the `indices` based on `cache_position >= self.config.attention_window_size - 1`.
+ The `to_shift` is only true once we are above attention_window_size. Thus with `attention_window_size==64`:
+
+ indices = (slicing + to_shift[-1].int()-1) % self.config.attention_window_size
+ tensor([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
+ 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
+ 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54,
+ 55, 56, 57, 58, 59, 60, 61, 62, 63, 0])
+
+ We overwrite the cache using these, then we always write at cache_position (clamped to `attention_window_size`)
+ """
+ cache_position = cache_kwargs.get("cache_position")
+ if cache_position.shape[0] > self.config.attention_window_size:
+ # int indexing -> device sync? in compile, use tensor
+ k_out = key_states[:, :, -self.config.attention_window_size :, :]
+ v_out = value_states[:, :, -self.config.attention_window_size :, :]
+ else:
+ slicing = torch.ones(
+ self.config.attention_window_size, dtype=torch.long, device=value_states.device
+ ).cumsum(0)
+ cache_position = cache_position.clamp(0, self.config.attention_window_size - 1)
+ to_shift = cache_position >= self.config.attention_window_size - 1
+ indices = (slicing + to_shift[-1].int() - 1) % self.config.attention_window_size
+
+ k_out, v_out = self.key_states.to(key_states.device), self.value_states.to(value_states.device)
+ k_out = k_out[:, :, indices]
+ v_out = v_out[:, :, indices]
+
+ k_out[:, :, cache_position] = key_states.to(k_out.dtype)
+ v_out[:, :, cache_position] = value_states.to(v_out.dtype)
+
+ self.key_states, self.value_states = k_out, v_out
+ return k_out, v_out
+
+
+class SqrtBoundDerivative(torch.autograd.Function):
+ """Computes a square root with a gradient clipped at `_MAX_SQRT_GRADIENT`."""
+
+ @staticmethod
+ def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+ """The forward pass, which is a normal `sqrt`."""
+ ctx.save_for_backward(x)
+ return torch.sqrt(x)
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+ """The backward pass, which clips the `sqrt` gradient."""
+ (x,) = ctx.saved_tensors
+ clipped_x_times_4 = torch.clip(4.0 * x, min=1 / (_MAX_SQRT_GRADIENT**2))
+ return grad_output / torch.sqrt(clipped_x_times_4)
+
+
+class RecurrentGemmaRglru(nn.Module):
+ """A Real-Gated Linear Recurrent Unit (RG-LRU) layer."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.num_attention_heads = config.num_attention_heads
+ self.block_width = config.lru_width // self.num_attention_heads
+
+ self.recurrent_param = nn.Parameter(torch.empty([config.lru_width]))
+ self.input_gate_weight = nn.Parameter(
+ torch.empty([self.num_attention_heads, self.block_width, self.block_width])
+ )
+ self.input_gate_bias = nn.Parameter(torch.empty([self.num_attention_heads, self.block_width]))
+
+ self.recurrent_gate_weight = nn.Parameter(
+ torch.empty([self.num_attention_heads, self.block_width, self.block_width])
+ )
+ self.recurrent_gate_bias = nn.Parameter(torch.empty([self.num_attention_heads, self.block_width]))
+ self.recurrent_states = None
+
+ def forward(
+ self,
+ activations: torch.Tensor,
+ position_ids: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ batch_size, seq_len, lru_width = activations.shape
+ reset = position_ids[:, :, None] == 0
+
+ reshape_act = activations.reshape(batch_size * seq_len, self.num_attention_heads, self.block_width)
+ reshape_act = reshape_act.permute(1, 0, 2)
+
+ res = torch.baddbmm(self.input_gate_bias[:, None, :], reshape_act, self.input_gate_weight)
+ input_gate = torch.sigmoid(res.transpose(0, 1).reshape(batch_size, seq_len, lru_width))
+
+ res = torch.baddbmm(self.recurrent_gate_bias[:, None, :], reshape_act, self.recurrent_gate_weight)
+ recurrent_gate = torch.sigmoid(res.transpose(0, 1).reshape(batch_size, seq_len, lru_width))
+
+ # Compute the parameter `A` of the recurrence.
+ log_recurrent_gate = -8.0 * recurrent_gate * nn.functional.softplus(self.recurrent_param)
+ recurrent_gate = torch.exp(log_recurrent_gate)
+ a_square = torch.exp(2 * log_recurrent_gate)
+
+ # Gate the input.
+ gated_inputs = activations * input_gate
+
+ # Apply gamma normalization to the input. We need to clip the derivatives of
+ # `sqrt` in order to prevent NaNs during training in bfloat16. TODO a bit annoying
+ multiplier = 1
+ tracing = isinstance(activations, torch.fx.Proxy) or is_torchdynamo_compiling()
+ if not torch.jit.is_tracing() and not tracing:
+ multiplier = SqrtBoundDerivative.apply(1 - a_square)
+ multiplier = reset + ~reset * multiplier
+ normalized_x = gated_inputs * multiplier.type(activations.dtype)
+
+ hidden_states, recurrent_states = self._rnn_scan(
+ hidden_states=normalized_x,
+ recurrent_gate=recurrent_gate,
+ reset=reset,
+ recurrent_states=self.recurrent_states,
+ )
+ self.recurrent_states = recurrent_states
+ return hidden_states
+
+ # TODO refactor
+ def _rnn_scan(
+ self,
+ hidden_states: torch.Tensor,
+ recurrent_gate: torch.Tensor,
+ reset: torch.Tensor,
+ recurrent_states: Union[torch.Tensor, None],
+ acc_dtype: torch.dtype = torch.float32,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Runs the recurrence of a linear RNN.
+
+ Args:
+ hidden_states: The input sequence.
+ recurrent_gate: The diagonal of the recurrence matrix `A`.
+ reset: Indicator of document boundaries, e.g. when to reset the hidden state
+ of the RNN.
+ recurrent_states: The initial hidden state.
+ acc_dtype: The data type for the accumulation.
+
+ Returns:
+ The output of the linear recurrence.
+ """
+ # Multiply `a` by the reset.
+ recurrent_gate = recurrent_gate * ~reset
+
+ if hidden_states.shape[1] == 1:
+ # Using scan in sampling mode.
+ if recurrent_states is None: # same here, when decoding you always have cache
+ return hidden_states, hidden_states[:, 0].type(acc_dtype)
+
+ else:
+ contextualized_states = recurrent_gate.type(acc_dtype) * recurrent_states[:, None].to(
+ recurrent_gate.device
+ )
+ contextualized_states += hidden_states.type(acc_dtype)
+ return contextualized_states.type(hidden_states.dtype), contextualized_states[:, -1]
+
+ else:
+ # Using scan in linear mode.
+ if recurrent_states is None:
+ recurrent_states = torch.zeros(hidden_states[:, 0].shape, dtype=acc_dtype, device=hidden_states.device)
+
+ contextualized_states = torch.zeros_like(hidden_states)
+ for t in range(hidden_states.shape[1]):
+ recurrent_states = recurrent_gate[:, t].type(acc_dtype) * recurrent_states.to(recurrent_gate.device)
+ recurrent_states = recurrent_states + hidden_states[:, t].type(acc_dtype)
+ contextualized_states[:, t] = recurrent_states.type(hidden_states.dtype)
+
+ return contextualized_states, recurrent_states
+
+
+class RecurrentGemmaRecurrentBlock(nn.Module):
+ """Griffin and Hawk's recurrent block."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.lru_width = config.lru_width
+ self.hidden_size = config.hidden_size
+ self.linear_y = nn.Linear(in_features=config.hidden_size, out_features=config.lru_width)
+ self.linear_x = nn.Linear(in_features=config.hidden_size, out_features=config.lru_width)
+ self.linear_out = nn.Linear(in_features=config.lru_width, out_features=config.hidden_size)
+ self.conv1d_width = config.conv1d_width
+ self.conv_1d = nn.Conv1d(
+ config.lru_width,
+ config.lru_width,
+ kernel_size=config.conv1d_width,
+ groups=config.lru_width,
+ padding=config.conv1d_width - 1,
+ )
+ self.rg_lru = RecurrentGemmaRglru(config)
+ self.act_fn = ACT2FN[config.hidden_activation]
+
+ self.conv1d_state = None
+
+ def forward(
+ self,
+ input_states: torch.Tensor,
+ position_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ cache_position: torch.Tensor,
+ use_cache: bool = True,
+ ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+ _, seq_len, _ = input_states.shape
+
+ y_branch = self.linear_y(input_states)
+ y_branch = self.act_fn(y_branch)
+
+ x_branch = self.linear_x(input_states)
+ x_branch = x_branch.transpose(1, 2)
+
+ if use_cache:
+ if cache_position.shape[0] != 1: # prefill
+ self.conv1d_state = nn.functional.pad(x_branch, (self.conv1d_width - x_branch.shape[-1] - 1, 0))
+ x_branch = self.conv_1d(x_branch)[..., :seq_len]
+ else: # decoding
+ conv_state = torch.cat((self.conv1d_state, x_branch), -1)
+ x_branch = torch.sum(conv_state * self.conv_1d.weight[:, 0, :], dim=-1) + self.conv_1d.bias
+ x_branch = x_branch.unsqueeze(-1)
+ self.conv1d_state = conv_state[:, :, 1:]
+ else:
+ x_branch = self.conv_1d(x_branch)[..., :seq_len]
+
+ x_branch = self.rg_lru(x_branch.transpose(1, 2), position_ids)
+
+ hidden_states = x_branch * y_branch
+ hidden_states = self.linear_out(hidden_states)
+ return hidden_states
+
+ def _setup_cache(self, batch, device, dtype):
+ # recurrent_states always computed in full precision
+ self.rg_lru.recurrent_states = torch.zeros((batch, self.lru_width), device=device, dtype=torch.float32)
+ self.conv1d_state = torch.zeros((batch, self.hidden_size, self.conv1d_width - 1), device=device, dtype=dtype)
+
+
+TEMPORAL_BLOCK_CLASSES = {"recurrent": RecurrentGemmaRecurrentBlock, "attention": RecurrentGemmaSdpaAttention}
+
+
+class RecurrentGemmaMlp(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size // 2
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.hidden_activation]
+
+ def forward(self, hidden_states):
+ gate = self.act_fn(self.gate_proj(hidden_states))
+ return self.down_proj(gate * self.up_proj(hidden_states))
+
+
+class RecurrentGemmaDecoderLayer(nn.Module):
+ """Griffin and Hawk's residual block."""
+
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.temporal_pre_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.temporal_block = TEMPORAL_BLOCK_CLASSES[config.layers_block_type[layer_idx]](config)
+ self.channel_pre_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.mlp_block = RecurrentGemmaMlp(config)
+
+ def forward(
+ self,
+ activations: torch.Tensor,
+ position_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ cache_position: Optional[torch.Tensor] = None,
+ use_cache: Optional[bool] = None,
+ ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+ raw_activations = activations
+ inputs_normalized = self.temporal_pre_norm(raw_activations) # RMSNorm introduces slight slight differences
+
+ hidden_states = self.temporal_block(
+ inputs_normalized, position_ids, attention_mask, cache_position=cache_position, use_cache=use_cache
+ )
+
+ residual = hidden_states + raw_activations
+
+ hidden_states = self.channel_pre_norm(residual)
+ hidden_states = self.mlp_block(hidden_states)
+
+ hidden_states = hidden_states + residual
+ return hidden_states
+
+
+RECURRENTGEMMA_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 ([`RecurrentGemmaConfig`]):
+ 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 RecurrentGemma Model outputting raw hidden-states without any specific head on top.",
+ RECURRENTGEMMA_START_DOCSTRING,
+)
+class RecurrentGemmaPreTrainedModel(PreTrainedModel):
+ config_class = RecurrentGemmaConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["RecurrentGemmaDecoderLayer"]
+ _skip_keys_device_placement = ["cache"]
+ _supports_flash_attn_2 = False
+ _supports_sdpa = False # we can't compare with eager for now
+ _supports_cache_class = True
+ _supports_quantized_cache = True
+
+ def _init_weights(self, module):
+ std = math.sqrt(self.config.w_init_variance_scale / self.config.conv1d_width)
+ if isinstance(module, nn.Conv1d):
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+ torch.nn.init.zeros_(module.bias)
+ elif isinstance(module, RecurrentGemmaSdpaAttention):
+ torch.nn.init.normal_(module.q_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size))
+ torch.nn.init.normal_(module.k_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size))
+ torch.nn.init.normal_(module.v_proj.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size))
+
+ std = math.sqrt(self.config.final_w_init_variance_scale / self.config.hidden_size)
+ torch.nn.init.normal_(module.o_proj.weight, mean=0.0, std=std)
+ elif isinstance(module, RecurrentGemmaRecurrentBlock):
+ torch.nn.init.zeros_(module.linear_x.bias)
+ torch.nn.init.normal_(module.linear_x.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size))
+
+ torch.nn.init.zeros_(module.linear_y.bias)
+ torch.nn.init.normal_(module.linear_y.weight, mean=0.0, std=math.sqrt(1.0 / self.config.hidden_size))
+
+ std = math.sqrt(self.config.final_w_init_variance_scale / self.config.lru_width)
+ torch.nn.init.normal_(module.linear_out.weight, mean=0.0, std=std)
+ torch.nn.init.zeros_(module.linear_out.bias)
+ elif isinstance(module, RecurrentGemmaRglru):
+ std = math.sqrt(
+ self.config.w_init_variance_scale / (self.config.lru_width // self.config.num_attention_heads)
+ )
+ torch.nn.init.normal_(module.input_gate_weight, mean=0.0, std=std)
+ torch.nn.init.normal_(module.recurrent_gate_weight, mean=0.0, std=std)
+ torch.nn.init.zeros_(module.input_gate_bias)
+ torch.nn.init.zeros_(module.recurrent_gate_bias)
+
+ module.recurrent_param.data.uniform_(0.9**2 + 1e-8, 0.999**2 + 1e-8)
+ module.recurrent_param.data.log_().mul_(0.5)
+ module.recurrent_param.data.neg_().exp_().sub_(1.0).log_()
+ elif isinstance(module, nn.Linear):
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+ if getattr(module, "bias", None) is not None:
+ torch.nn.init.zeros_(module.bias)
+
+ def _setup_cache(self, config, batch, device, dtype):
+ layers = getattr(self, "model", self).layers
+ for layer in layers:
+ layer.temporal_block._setup_cache(batch, device, dtype)
+
+ def reset_cache(self, batch, device, dtype):
+ pass
+
+
+RECURRENTGEMMA_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.
+
+ 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)
+ 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_hidden_states (`bool`, *optional*):
+ Whether or not to return the hidden states of all 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 RecurrentGemma Model outputting raw hidden-states without any specific head on top.",
+ RECURRENTGEMMA_START_DOCSTRING,
+)
+class RecurrentGemmaModel(RecurrentGemmaPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`RecurrentGemmaDecoderLayer`]
+
+ Args:
+ config: RecurrentGemmaConfig
+ """
+
+ def __init__(self, config: RecurrentGemmaConfig):
+ 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(
+ [RecurrentGemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.final_norm = RecurrentGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.gradient_checkpointing = False
+
+ self.register_buffer(
+ "normalizer", torch.tensor(self.config.hidden_size**0.5, dtype=torch.bfloat16), persistent=False
+ )
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel.get_input_embeddings
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel.set_input_embeddings
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ @add_start_docstrings_to_model_forward(RECURRENTGEMMA_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithNoAttention]:
+ 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 must specify exactly one of input_ids or inputs_embeds")
+
+ 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)
+
+ hidden_states = inputs_embeds
+
+ if use_cache and inputs_embeds.shape[1] != 1: # TODO let's maybe only call in the `generate`?
+ self._setup_cache(self.config, hidden_states.shape[0], hidden_states.device, hidden_states.dtype)
+
+ if cache_position is None:
+ cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
+
+ hidden_states = hidden_states * self.normalizer.type(hidden_states.dtype)
+
+ all_hidden_states = () if output_hidden_states else None
+ for i, residual_block in enumerate(self.layers):
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+ if self.gradient_checkpointing and self.training:
+ hidden_states = self._gradient_checkpointing_func(
+ residual_block.__call__, hidden_states, position_ids, causal_mask, cache_position, use_cache
+ )
+ else:
+ hidden_states = residual_block(hidden_states, position_ids, causal_mask, cache_position, use_cache)
+
+ hidden_states = self.final_norm(hidden_states)
+
+ # 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, all_hidden_states] if v is not None)
+
+ return BaseModelOutputWithNoAttention(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ )
+
+ # Ignore copy
+ def _update_causal_mask(self, attention_mask, input_tensor, cache_position):
+ dtype, device = input_tensor.dtype, input_tensor.device
+ min_dtype = torch.finfo(dtype).min
+ sequence_length = input_tensor.shape[1]
+ target_length = max(self.config.attention_window_size, sequence_length)
+
+ diagonal = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
+ causal_mask = diagonal
+ if sequence_length != 1:
+ causal_mask = torch.triu(diagonal, 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)
+
+ if attention_mask is not None and attention_mask.device.type in ["cuda", "xpu"]:
+ # 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
+
+
+# TODO: re-enable check: Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->RECURRENTGEMMA,Llama->RecurrentGemma,llama->gemma
+class RecurrentGemmaForCausalLM(RecurrentGemmaPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = RecurrentGemmaModel(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(RECURRENTGEMMA_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=CausalLMOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ use_cache: Optional[bool] = None,
+ **kwargs,
+ ) -> Union[Tuple, CausalLMOutput]:
+ 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:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, RecurrentGemmaForCausalLM
+
+ >>> model = RecurrentGemmaForCausalLM.from_pretrained("google/recurrentgemma-2b")
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/recurrentgemma-2b")
+
+ >>> 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_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
+ output_hidden_states = True
+ outputs = self.model(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ cache_position=cache_position,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+
+ # Soft-cap the logits TODO remove if always done.
+ # if self.config.logits_soft_cap is not None:
+ cap = self.config.logits_soft_cap
+ logits = nn.functional.tanh(logits / cap) * cap
+
+ loss = None
+ if labels is not None:
+ # Upcast to float if we need to compute the loss to avoid potential precision issues
+ logits = logits.float()
+ loss = self.loss_function(
+ logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return CausalLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ )
+
+ # Ignore copy
+ def _reorder_cache(self, past_key_values, beam_idx):
+ for layer in self.layers:
+ if hasattr(layer.temporal_block, "key_states"):
+ k_state = layer.temporal_block.key_states
+ v_state = layer.temporal_block.value_states
+ k_state = k_state.index_select(0, beam_idx.to(k_state.device))
+ v_state = v_state.index_select(0, beam_idx.to(v_state.device))
+ return None
+
+
+__all__ = ["RecurrentGemmaForCausalLM", "RecurrentGemmaModel", "RecurrentGemmaPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/swin/__init__.py b/lib/python3.12/site-packages/transformers/models/swin/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3dc5871b0375e8a5553e30dce1e5dde807a1b493
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swin/__init__.py
@@ -0,0 +1,28 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_swin import *
+ from .modeling_swin import *
+ from .modeling_tf_swin import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/swin/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dc038c60239f345e894f486c0d535b7a6d786917
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swin/__pycache__/configuration_swin.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/configuration_swin.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4b51f0a2e14532abeda8c753f35decf74874e8f2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/configuration_swin.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_swin.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_swin.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..22f371fe75e13e0d9f3d84f9dbb9056535eea696
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_swin.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_tf_swin.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_tf_swin.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..047f840570aee1cd20c0be943a0e1bd5bfa6db17
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swin/__pycache__/modeling_tf_swin.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py b/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..da6ba9871407a644b99435e69c6ec46e26fd8735
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swin/configuration_swin.py
@@ -0,0 +1,179 @@
+# 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.
+"""Swin Transformer 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__)
+
+
+class SwinConfig(BackboneConfigMixin, PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`SwinModel`]. It is used to instantiate a Swin
+ 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 Swin
+ [microsoft/swin-tiny-patch4-window7-224](https://huggingface.co/microsoft/swin-tiny-patch4-window7-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:
+ 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.
+ 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 SwinConfig, SwinModel
+
+ >>> # Initializing a Swin microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> configuration = SwinConfig()
+
+ >>> # Initializing a model (with random weights) from the microsoft/swin-tiny-patch4-window7-224 style configuration
+ >>> model = SwinModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "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,
+ 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.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
+ self.encoder_stride = encoder_stride
+ # 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))
+ 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
+ )
+
+
+class SwinOnnxConfig(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
+
+
+__all__ = ["SwinConfig", "SwinOnnxConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py b/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..5de428831e347dfe5d136f46f62574ad8c043393
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swin/modeling_swin.py
@@ -0,0 +1,1432 @@
+# 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 Swin Transformer model."""
+
+import collections.abc
+import math
+import warnings
+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, meshgrid, prune_linear_layer
+from ...utils import (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+ torch_int,
+)
+from ...utils.backbone_utils import BackboneMixin
+from .configuration_swin import SwinConfig
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "SwinConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "microsoft/swin-tiny-patch4-window7-224"
+_EXPECTED_OUTPUT_SHAPE = [1, 49, 768]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "microsoft/swin-tiny-patch4-window7-224"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+# drop_path, SwinPatchEmbeddings, SwinPatchMerging and SwinDropPath are from the timm library.
+
+
+@dataclass
+class SwinEncoderOutput(ModelOutput):
+ """
+ Swin 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: 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
+class SwinModelOutput(ModelOutput):
+ """
+ Swin 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: Optional[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
+class SwinMaskedImageModelingOutput(ModelOutput):
+ """
+ Swin 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.
+ 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
+ reconstruction: 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
+
+ @property
+ def logits(self):
+ warnings.warn(
+ "logits attribute is deprecated and will be removed in version 5 of Transformers."
+ " Please use the reconstruction attribute to retrieve the final output instead.",
+ FutureWarning,
+ )
+ return self.reconstruction
+
+
+@dataclass
+class SwinImageClassifierOutput(ModelOutput):
+ """
+ Swin 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: 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
+
+
+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
+
+
+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
+
+
+class SwinEmbeddings(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 = SwinPatchEmbeddings(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)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding
+ 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. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embeddings.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, :1]
+ patch_pos_embed = self.position_embeddings[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor],
+ bool_masked_pos: Optional[torch.BoolTensor] = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> Tuple[torch.Tensor]:
+ _, num_channels, height, width = pixel_values.shape
+ 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:
+ 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, output_dimensions
+
+
+class SwinPatchEmbeddings(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
+ # 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
+
+
+class SwinPatchMerging(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.beit.modeling_beit.BeitDropPath with Beit->Swin
+class SwinDropPath(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 SwinSelfAttention(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 SwinModel 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
+
+
+class SwinSelfOutput(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 SwinAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size):
+ super().__init__()
+ self.self = SwinSelfAttention(config, dim, num_heads, window_size)
+ self.output = SwinSelfOutput(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
+
+
+class SwinIntermediate(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
+
+
+class SwinOutput(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 SwinLayer(nn.Module):
+ def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, 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 = SwinAttention(config, dim, num_heads, window_size=self.window_size)
+ self.drop_path = SwinDropPath(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 = SwinIntermediate(config, dim)
+ self.output = SwinOutput(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 = torch_int(0)
+ self.window_size = (
+ torch.min(torch.tensor(input_resolution)) if torch.jit.is_tracing() else min(input_resolution)
+ )
+
+ def get_attn_mask(self, height, width, dtype, device):
+ if self.shift_size > 0:
+ # calculate attention mask for SW-MSA
+ img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)
+ 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, device=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
+
+
+class SwinStage(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(
+ [
+ SwinLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ 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
+
+
+class SwinEncoder(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(
+ [
+ SwinStage(
+ 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=SwinPatchMerging 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, SwinEncoderOutput]:
+ 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 SwinEncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ reshaped_hidden_states=all_reshaped_hidden_states,
+ )
+
+
+class SwinPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = SwinConfig
+ base_model_prefix = "swin"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["SwinStage"]
+
+ 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)
+ elif isinstance(module, SwinEmbeddings):
+ if module.mask_token is not None:
+ module.mask_token.data.zero_()
+ if module.position_embeddings is not None:
+ module.position_embeddings.data.zero_()
+ elif isinstance(module, SwinSelfAttention):
+ module.relative_position_bias_table.data.zero_()
+
+
+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 ([`SwinConfig`]): 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 [`ViTImageProcessor.__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.
+ interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare Swin Model transformer outputting raw hidden-states without any specific head on top.",
+ SWIN_START_DOCSTRING,
+ """
+ add_pooling_layer (`bool`, *optional*, defaults to `True`):
+ Whether or not to apply pooling layer.
+ use_mask_token (`bool`, *optional*, defaults to `False`):
+ Whether or not to create and apply mask tokens in the embedding layer.
+ """,
+)
+class SwinModel(SwinPreTrainedModel):
+ 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 = SwinEmbeddings(config, use_mask_token=use_mask_token)
+ self.encoder = SwinEncoder(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
+
+ 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=SwinModelOutput,
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, SwinModelOutput]:
+ 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, len(self.config.depths))
+
+ embedding_output, input_dimensions = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ 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]
+ 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 SwinModelOutput(
+ 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(
+ """Swin 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).
+
+
+ """,
+ SWIN_START_DOCSTRING,
+)
+class SwinForMaskedImageModeling(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.swin = SwinModel(config, add_pooling_layer=False, use_mask_token=True)
+
+ num_features = int(config.embed_dim * 2 ** (config.num_layers - 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(SWIN_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=SwinMaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, SwinMaskedImageModelingOutput]:
+ 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, SwinForMaskedImageModeling
+ >>> 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/swin-base-simmim-window6-192")
+ >>> model = SwinForMaskedImageModeling.from_pretrained("microsoft/swin-base-simmim-window6-192")
+
+ >>> 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, 192, 192]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.swin(
+ pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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 SwinMaskedImageModelingOutput(
+ loss=masked_im_loss,
+ reconstruction=reconstructed_pixel_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@add_start_docstrings(
+ """
+ Swin 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.
+
+
+
+ Note that it's possible to fine-tune Swin on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """,
+ SWIN_START_DOCSTRING,
+)
+class SwinForImageClassification(SwinPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.swin = SwinModel(config)
+
+ # Classifier head
+ self.classifier = (
+ nn.Linear(self.swin.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(SWIN_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_IMAGE_CLASS_CHECKPOINT,
+ output_type=SwinImageClassifierOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
+ )
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, SwinImageClassifierOutput]:
+ 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.swin(
+ pixel_values,
+ 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 = 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 SwinImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@add_start_docstrings(
+ """
+ Swin backbone, to be used with frameworks like DETR and MaskFormer.
+ """,
+ SWIN_START_DOCSTRING,
+)
+class SwinBackbone(SwinPreTrainedModel, BackboneMixin):
+ def __init__(self, config: SwinConfig):
+ super().__init__(config)
+ super()._init_backbone(config)
+
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+ self.embeddings = SwinEmbeddings(config)
+ self.encoder = SwinEncoder(config, self.embeddings.patch_grid)
+
+ # 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
+
+ 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(
+ ... "microsoft/swin-tiny-patch4-window7-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, 768, 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, input_dimensions = self.embeddings(pixel_values)
+
+ outputs = self.encoder(
+ embedding_output,
+ input_dimensions,
+ head_mask=None,
+ output_attentions=output_attentions,
+ output_hidden_states=True,
+ output_hidden_states_before_downsampling=True,
+ always_partition=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,
+ )
+
+
+__all__ = [
+ "SwinForImageClassification",
+ "SwinForMaskedImageModeling",
+ "SwinModel",
+ "SwinPreTrainedModel",
+ "SwinBackbone",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/swin/modeling_tf_swin.py b/lib/python3.12/site-packages/transformers/models/swin/modeling_tf_swin.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1c4fb1620fd2a0feb4a352b1de12cb4e7fc6054
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swin/modeling_tf_swin.py
@@ -0,0 +1,1638 @@
+# 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 Swin Transformer model."""
+
+from __future__ import annotations
+
+import collections.abc
+import math
+import warnings
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+
+import tensorflow as tf
+
+from ...activations_tf import ACT2FN
+from ...modeling_tf_utils import (
+ TFPreTrainedModel,
+ TFSequenceClassificationLoss,
+ get_initializer,
+ keras,
+ keras_serializable,
+ unpack_inputs,
+)
+from ...tf_utils import shape_list
+from ...utils import (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from .configuration_swin import SwinConfig
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "SwinConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "microsoft/swin-tiny-patch4-window7-224"
+_EXPECTED_OUTPUT_SHAPE = [1, 49, 768]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "microsoft/swin-tiny-patch4-window7-224"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+
+# drop_path, TFSwinPatchEmbeddings, TFSwinPatchMerging and TFSwinDropPath are tensorflow
+# implementations of PyTorch functionalities in the timm library.
+
+
+@dataclass
+class TFSwinEncoderOutput(ModelOutput):
+ """
+ Swin encoder's outputs, with potential hidden states and attentions.
+
+ 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.
+ 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 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tf.Tensor` (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(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 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: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ reshaped_hidden_states: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFSwinModelOutput(ModelOutput):
+ """
+ Swin model's outputs that also contains a pooling of the last hidden states.
+
+ 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)`, *optional*, returned when `add_pooling_layer=True` is passed):
+ Average pooling of the last layer hidden-state.
+ 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 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tf.Tensor` (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(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 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: Optional[tf.Tensor] = None
+ pooler_output: tf.Tensor | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ reshaped_hidden_states: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFSwinMaskedImageModelingOutput(ModelOutput):
+ """
+ Swin masked image model outputs.
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
+ Masked image modeling (MLM) loss.
+ reconstruction (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
+ Reconstructed pixel values.
+ 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 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tf.Tensor` (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(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 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: tf.Tensor | None = None
+ reconstruction: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ reshaped_hidden_states: Tuple[tf.Tensor, ...] | None = None
+
+ @property
+ def logits(self):
+ warnings.warn(
+ "logits attribute is deprecated and will be removed in version 5 of Transformers."
+ " Please use the reconstruction attribute to retrieve the final output instead.",
+ FutureWarning,
+ )
+ return self.reconstruction
+
+
+@dataclass
+class TFSwinImageClassifierOutput(ModelOutput):
+ """
+ Swin outputs for image classification.
+
+ 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).
+ 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 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tf.Tensor` (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(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 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: tf.Tensor | None = None
+ logits: Optional[tf.Tensor] = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+ reshaped_hidden_states: Tuple[tf.Tensor, ...] | None = None
+
+
+def window_partition(input_feature: tf.Tensor, window_size: int) -> tf.Tensor:
+ """
+ Partitions the given input into windows.
+ """
+ batch_size, height, width, num_channels = shape_list(input_feature)
+ input_feature = tf.reshape(
+ input_feature,
+ (batch_size, height // window_size, window_size, width // window_size, window_size, num_channels),
+ )
+ windows = tf.transpose(input_feature, (0, 1, 3, 2, 4, 5))
+ windows = tf.reshape(windows, (-1, window_size, window_size, num_channels))
+ return windows
+
+
+def window_reverse(windows: tf.Tensor, window_size: int, height: int, width: int) -> tf.Tensor:
+ """
+ Merges windows to produce higher resolution features.
+ """
+ x = tf.shape(windows)[0]
+ y = tf.cast(height * width / (window_size * window_size), tf.int32)
+ batch_size = tf.math.floordiv(x, y)
+ windows = tf.reshape(
+ windows, (batch_size, height // window_size, width // window_size, window_size, window_size, -1)
+ )
+ windows = tf.transpose(windows, (0, 1, 3, 2, 4, 5))
+ windows = tf.reshape(windows, (batch_size, height, width, -1))
+ return windows
+
+
+def drop_path(
+ input: tf.Tensor, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+) -> tf.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ input_shape = shape_list(input)
+ ndim = len(input_shape)
+ shape = [input_shape[0]] + [1] * (ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = tf.random.uniform(shape)
+ random_tensor = tf.where(random_tensor <= keep_prob, 1.0, 0.0)
+ if keep_prob > 0.0 and scale_by_keep:
+ random_tensor /= keep_prob
+ return input * random_tensor
+
+
+class TFSwinEmbeddings(keras.layers.Layer):
+ """
+ Construct the patch and position embeddings. Optionally, also the mask token.
+ """
+
+ def __init__(self, config: SwinConfig, use_mask_token: bool = False, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.patch_embeddings = TFSwinPatchEmbeddings(config, name="patch_embeddings")
+ self.num_patches = self.patch_embeddings.num_patches
+ self.patch_grid = self.patch_embeddings.grid_size
+ self.embed_dim = config.embed_dim
+ self.use_mask_token = use_mask_token
+ self.use_absolute_embeddings = config.use_absolute_embeddings
+
+ self.norm = keras.layers.LayerNormalization(name="norm", epsilon=1e-5)
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob, name="dropout")
+ self.config = config
+
+ def build(self, input_shape: tf.TensorShape) -> None:
+ if self.use_mask_token:
+ self.mask_token = self.add_weight(shape=(1, 1, self.embed_dim), initializer="zeros", name="mask_token")
+ else:
+ self.mask_token = None
+
+ if self.use_absolute_embeddings:
+ self.position_embeddings = self.add_weight(
+ (1, self.num_patches + 1, self.embed_dim), initializer="zeros", name="positional_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)
+ if getattr(self, "norm", None) is not None:
+ with tf.name_scope(self.norm.name):
+ self.norm.build([None, None, self.config.embed_dim])
+ 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: Optional[bool] = None, training: bool = False
+ ) -> Tuple[tf.Tensor, Tuple[int, int]]:
+ embeddings, output_dimensions = self.patch_embeddings(pixel_values, training=training)
+ embeddings = self.norm(embeddings, training=training)
+ batch_size, seq_len, _ = shape_list(embeddings)
+
+ if bool_masked_pos is not None:
+ mask_tokens = tf.repeat(self.mask_token, batch_size, 0)
+ mask_tokens = tf.repeat(mask_tokens, seq_len, 1)
+ # replace the masked visual tokens by mask_tokens
+ mask = tf.expand_dims(bool_masked_pos, -1)
+ mask = tf.cast(mask, mask_tokens.dtype)
+
+ embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
+
+ if self.position_embeddings is not None:
+ embeddings = embeddings + self.position_embeddings
+
+ embeddings = self.dropout(embeddings, training=training)
+
+ return embeddings, output_dimensions
+
+
+class TFSwinPatchEmbeddings(keras.layers.Layer):
+ """
+ Image to Patch Embedding.
+ """
+
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+ 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 = keras.layers.Conv2D(
+ filters=hidden_size,
+ kernel_size=self.patch_size,
+ strides=self.patch_size,
+ padding="valid",
+ name="projection",
+ )
+
+ def maybe_pad(self, pixel_values: tf.Tensor, height: int, width: int) -> tf.Tensor:
+ if width % self.patch_size[1] != 0:
+ pad_values = ((0, 0), (0, 0), (0, 0), (0, self.patch_size[1] - width % self.patch_size[1]))
+ pixel_values = tf.pad(pixel_values, pad_values)
+ if height % self.patch_size[0] != 0:
+ pad_values = ((0, 0), (0, 0), (0, self.patch_size[0] - height % self.patch_size[0]), (0, 0))
+ pixel_values = tf.pad(pixel_values, pad_values)
+ return pixel_values
+
+ def call(self, pixel_values: tf.Tensor, training: bool = False) -> Tuple[tf.Tensor, Tuple[int, int]]:
+ _, num_channels, height, width = 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."
+ )
+ # pad the input to be divisible by self.patch_size, if needed
+ pixel_values = self.maybe_pad(pixel_values, height, width)
+
+ # B,C,H,W -> B,H,W,C
+ pixel_values = tf.transpose(pixel_values, (0, 2, 3, 1))
+
+ embeddings = self.projection(pixel_values, training=training)
+
+ # B,H,W,C -> B,C,H,W
+ embeddings = tf.transpose(embeddings, (0, 3, 1, 2))
+
+ batch_size, channels, height, width = shape_list(embeddings)
+ output_dimensions = (height, width)
+
+ embeddings = tf.reshape(embeddings, (batch_size, channels, -1))
+ embeddings = tf.transpose(embeddings, (0, 2, 1))
+ return embeddings, output_dimensions
+
+ 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 TFSwinPatchMerging(keras.layers.Layer):
+ """
+ Patch Merging Layer.
+
+ Args:
+ input_resolution (`Tuple[int]`):
+ Resolution of input feature.
+ dim (`int`):
+ Number of input channels.
+ norm_layer (`keras.layer.Layer`, *optional*, defaults to `keras.layers.LayerNormalization`):
+ Normalization layer class.
+ """
+
+ def __init__(
+ self, input_resolution: Tuple[int, int], dim: int, norm_layer: Optional[Callable] = None, **kwargs
+ ) -> None:
+ super().__init__(**kwargs)
+ self.input_resolution = input_resolution
+ self.dim = dim
+ self.reduction = keras.layers.Dense(2 * dim, use_bias=False, name="reduction")
+ if norm_layer is None:
+ # Use same default epsilon as PyTorch
+ self.norm = keras.layers.LayerNormalization(epsilon=1e-5, name="norm")
+ else:
+ self.norm = norm_layer(name="norm")
+
+ def maybe_pad(self, input_feature: tf.Tensor, height: int, width: int) -> tf.Tensor:
+ should_pad = (height % 2 == 1) or (width % 2 == 1)
+ if should_pad:
+ pad_values = ((0, 0), (0, height % 2), (0, width % 2), (0, 0))
+ input_feature = tf.pad(input_feature, pad_values)
+
+ return input_feature
+
+ def call(self, input_feature: tf.Tensor, input_dimensions: Tuple[int, int], training: bool = False) -> tf.Tensor:
+ height, width = input_dimensions
+ # `dim` is height * width
+ batch_size, _, num_channels = shape_list(input_feature)
+
+ input_feature = tf.reshape(input_feature, (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 = tf.concat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
+ input_feature = tf.reshape(
+ input_feature, (batch_size, -1, 4 * num_channels)
+ ) # batch_size height/2*width/2 4*C
+
+ input_feature = self.norm(input_feature, training=training)
+ input_feature = self.reduction(input_feature, training=training)
+
+ return input_feature
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "reduction", None) is not None:
+ with tf.name_scope(self.reduction.name):
+ self.reduction.build([None, None, 4 * self.dim])
+ if getattr(self, "norm", None) is not None:
+ with tf.name_scope(self.norm.name):
+ self.norm.build([None, None, 4 * self.dim])
+
+
+class TFSwinDropPath(keras.layers.Layer):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: Optional[float] = None, scale_by_keep: bool = True, **kwargs) -> None:
+ super(TFSwinDropPath, self).__init__(**kwargs)
+ self.drop_prob = drop_prob
+ self.scale_by_keep = scale_by_keep
+
+ def call(self, input: tf.Tensor, training: bool = False) -> tf.Tensor:
+ return drop_path(input, self.drop_prob, training, self.scale_by_keep)
+
+
+class TFSwinSelfAttention(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, dim: int, num_heads: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ 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
+ window_size = config.window_size
+ self.window_size = (
+ window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
+ )
+
+ self.query = keras.layers.Dense(
+ self.all_head_size,
+ kernel_initializer=get_initializer(config.initializer_range),
+ use_bias=config.qkv_bias,
+ name="query",
+ )
+ self.key = keras.layers.Dense(
+ self.all_head_size,
+ kernel_initializer=get_initializer(config.initializer_range),
+ use_bias=config.qkv_bias,
+ name="key",
+ )
+ self.value = keras.layers.Dense(
+ self.all_head_size,
+ kernel_initializer=get_initializer(config.initializer_range),
+ use_bias=config.qkv_bias,
+ name="value",
+ )
+
+ self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
+
+ def build(self, input_shape: tf.TensorShape) -> None:
+ self.relative_position_bias_table = self.add_weight(
+ shape=(((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1)), self.num_attention_heads),
+ initializer="zeros",
+ name="relative_position_bias_table",
+ )
+ self.relative_position_index = self.add_weight(
+ shape=(self.window_size[0] ** 2, self.window_size[1] ** 2),
+ trainable=False,
+ dtype=tf.int32,
+ name="relative_position_index",
+ )
+
+ # get pair-wise relative position index for each token inside the window
+ coords_h = tf.range(self.window_size[0])
+ coords_w = tf.range(self.window_size[1])
+ coords = tf.stack(tf.meshgrid(coords_h, coords_w, indexing="ij"))
+ coords_flatten = tf.reshape(coords, (shape_list(coords)[0], -1))
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+ relative_coords = tf.transpose(relative_coords, (1, 2, 0))
+
+ stack_0, stack_1 = tf.unstack(relative_coords, axis=2)
+ stack_0 += self.window_size[0] - 1
+ stack_0 *= 2 * self.window_size[1] - 1
+ stack_1 += self.window_size[1] - 1
+ relative_coords = tf.stack([stack_0, stack_1], axis=2)
+
+ self.relative_position_index.assign(tf.cast(tf.reduce_sum(relative_coords, axis=-1), tf.int32))
+
+ 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.all_head_size])
+ if getattr(self, "key", None) is not None:
+ with tf.name_scope(self.key.name):
+ self.key.build([None, None, self.all_head_size])
+ if getattr(self, "value", None) is not None:
+ with tf.name_scope(self.value.name):
+ self.value.build([None, None, self.all_head_size])
+
+ 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, (0, 2, 1, 3))
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ attention_mask: tf.Tensor | None = None,
+ head_mask: tf.Tensor | None = None,
+ output_attentions: bool = False,
+ training: bool = False,
+ ) -> Tuple[tf.Tensor, ...]:
+ batch_size, dim, _ = shape_list(hidden_states)
+ 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 = tf.matmul(query_layer, tf.transpose(key_layer, (0, 1, 3, 2)))
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ relative_position_bias = tf.gather(
+ self.relative_position_bias_table, tf.reshape(self.relative_position_index, (-1,))
+ )
+ relative_position_bias = tf.reshape(
+ relative_position_bias,
+ (self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1),
+ )
+
+ relative_position_bias = tf.transpose(relative_position_bias, (2, 0, 1))
+ attention_scores = attention_scores + tf.expand_dims(relative_position_bias, 0)
+
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in SwinModel call() function)
+ mask_shape = shape_list(attention_mask)[0]
+ attention_scores = tf.reshape(
+ attention_scores, (batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim)
+ )
+ attention_mask = tf.expand_dims(attention_mask, 1)
+ attention_mask = tf.expand_dims(attention_mask, 0)
+ attention_scores = attention_scores + attention_mask
+ attention_scores = tf.reshape(attention_scores, (-1, self.num_attention_heads, dim, dim))
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = tf.nn.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 = tf.matmul(attention_probs, value_layer)
+ context_layer = tf.transpose(context_layer, (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,)
+
+ return outputs
+
+
+class TFSwinSelfOutput(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, dim: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(dim, name="dense")
+ self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob, name="dropout")
+ self.dim = dim
+
+ def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(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.dim])
+ if getattr(self, "dropout", None) is not None:
+ with tf.name_scope(self.dropout.name):
+ self.dropout.build(None)
+
+
+class TFSwinAttention(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, dim: int, num_heads: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.self = TFSwinSelfAttention(config, dim, num_heads, name="self")
+ self.self_output = TFSwinSelfOutput(config, dim, name="output")
+ self.pruned_heads = set()
+
+ def prune_heads(self, heads):
+ """
+ Prunes heads of the model. See base class PreTrainedModel heads: dict of {layer_num: list of heads to prune in
+ this layer}
+ """
+ raise NotImplementedError
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ attention_mask: tf.Tensor | None = None,
+ head_mask: tf.Tensor | None = None,
+ output_attentions: bool = False,
+ training: bool = False,
+ ) -> tf.Tensor:
+ self_outputs = self.self(hidden_states, attention_mask, head_mask, output_attentions, training=training)
+ attention_output = self.self_output(self_outputs[0], hidden_states, 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", None) is not None:
+ with tf.name_scope(self.self.name):
+ self.self.build(None)
+ if getattr(self, "self_output", None) is not None:
+ with tf.name_scope(self.self_output.name):
+ self.self_output.build(None)
+
+
+class TFSwinIntermediate(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, dim: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(int(config.mlp_ratio * dim), name="dense")
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+ self.dim = dim
+
+ def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
+ 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.dim])
+
+
+class TFSwinOutput(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, dim: int, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.dense = keras.layers.Dense(dim, name="dense")
+ self.dropout = keras.layers.Dropout(config.hidden_dropout_prob, "dropout")
+ self.config = config
+ self.dim = dim
+
+ def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(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, int(self.config.mlp_ratio * self.dim)])
+
+
+class TFSwinLayer(keras.layers.Layer):
+ def __init__(
+ self,
+ config,
+ dim,
+ input_resolution: Tuple[int, int],
+ num_heads: int,
+ drop_path_rate: float = 0.0,
+ shift_size: int = 0,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ min_res = tf.reduce_min(input_resolution)
+ self.window_size = min_res if min_res <= config.window_size else config.window_size
+ self.shift_size = 0 if min_res <= self.window_size else shift_size
+ self.input_resolution = input_resolution
+
+ self.layernorm_before = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_before")
+ self.attention = TFSwinAttention(config, dim, num_heads, name="attention")
+ self.drop_path = (
+ TFSwinDropPath(drop_path_rate, name="drop_path")
+ if drop_path_rate > 0.0
+ else keras.layers.Activation("linear", name="drop_path")
+ )
+ self.layernorm_after = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_after")
+ self.intermediate = TFSwinIntermediate(config, dim, name="intermediate")
+ self.swin_output = TFSwinOutput(config, dim, name="output")
+ self.dim = dim
+
+ def get_attn_mask(self, height: int, width: int, window_size: int, shift_size: int) -> tf.Tensor | None:
+ img_mask = tf.zeros((height, width))
+ height_slices = ((0, -window_size), (-window_size, -shift_size), (-shift_size, -1))
+ width_slices = ((0, -window_size), (-window_size, -shift_size), (-shift_size, -1))
+
+ # calculate attention mask for SW-MSA
+ if shift_size > 0:
+ count = 0
+ for height_slice in height_slices:
+ for width_slice in width_slices:
+ height_inds = tf.range(height_slice[0] % height, height_slice[1] % height + 1)
+ width_inds = tf.range(width_slice[0] % width, width_slice[1] % width + 1)
+ indices = tf.reshape(tf.stack(tf.meshgrid(height_inds, width_inds), axis=-1), (-1, 2))
+ if len(indices) >= 1:
+ updates = tf.ones((len(indices),), dtype=img_mask.dtype) * count
+ img_mask = tf.tensor_scatter_nd_update(img_mask, indices, updates)
+ count += 1
+
+ img_mask = tf.expand_dims(img_mask, -1)
+ img_mask = tf.expand_dims(img_mask, 0)
+
+ mask_windows = window_partition(img_mask, window_size)
+ mask_windows = tf.reshape(mask_windows, (-1, window_size * window_size))
+ attn_mask = tf.expand_dims(mask_windows, 1) - tf.expand_dims(mask_windows, 2)
+ attn_mask = tf.where(attn_mask != 0, float(-100.0), attn_mask)
+ attn_mask = tf.where(attn_mask == 0, float(0.0), attn_mask)
+ return attn_mask
+
+ def maybe_pad(
+ self, hidden_states: tf.Tensor, window_size: int, height: int, width: int
+ ) -> Tuple[tf.Tensor, tf.Tensor]:
+ pad_right = (window_size - width % window_size) % window_size
+ pad_bottom = (window_size - height % window_size) % window_size
+ pad_values = [[0, 0], [0, pad_bottom], [0, pad_right], [0, 0]]
+ hidden_states = tf.pad(hidden_states, pad_values)
+ pad_values = tf.reshape(pad_values, (-1,))
+ return hidden_states, pad_values
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ input_dimensions: Tuple[int, int],
+ head_mask: tf.Tensor | None = None,
+ output_attentions: bool = False,
+ training: bool = False,
+ ) -> tf.Tensor:
+ # if window size is larger than input resolution, we don't partition windows
+ min_res = tf.reduce_min(input_dimensions)
+ shift_size = 0 if min_res <= self.window_size else self.shift_size
+ window_size = min_res if min_res <= self.window_size else self.window_size
+
+ height, width = input_dimensions
+ batch_size, _, channels = shape_list(hidden_states)
+ shortcut = hidden_states
+
+ hidden_states = self.layernorm_before(hidden_states, training=training)
+ hidden_states = tf.reshape(hidden_states, (batch_size, height, width, channels))
+ # pad hidden_states to multiples of window size
+ hidden_states, pad_values = self.maybe_pad(hidden_states, window_size, height, width)
+
+ _, height_pad, width_pad, _ = shape_list(hidden_states)
+ # cyclic shift
+ if shift_size > 0:
+ shifted_hidden_states = tf.roll(hidden_states, shift=(-shift_size, -shift_size), axis=(1, 2))
+ else:
+ shifted_hidden_states = hidden_states
+
+ # partition windows
+ hidden_states_windows = window_partition(shifted_hidden_states, window_size)
+ hidden_states_windows = tf.reshape(hidden_states_windows, (-1, window_size * window_size, channels))
+ attn_mask = self.get_attn_mask(
+ height=height_pad, width=width_pad, window_size=window_size, shift_size=shift_size
+ )
+
+ attention_outputs = self.attention(
+ hidden_states_windows, attn_mask, head_mask, output_attentions=output_attentions, training=training
+ )
+
+ attention_output = attention_outputs[0]
+
+ attention_windows = tf.reshape(attention_output, (-1, window_size, window_size, channels))
+ shifted_windows = window_reverse(attention_windows, window_size, height_pad, width_pad)
+
+ # reverse cyclic shift
+ if shift_size > 0:
+ attention_windows = tf.roll(shifted_windows, shift=(shift_size, shift_size), axis=(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, :]
+
+ attention_windows = tf.reshape(attention_windows, (batch_size, height * width, channels))
+
+ hidden_states = shortcut + self.drop_path(attention_windows, training=training)
+
+ layer_output = self.layernorm_after(hidden_states, training=training)
+ layer_output = self.intermediate(layer_output)
+ layer_output = hidden_states + self.swin_output(layer_output, training=training)
+
+ layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
+ return layer_outputs
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "layernorm_before", None) is not None:
+ with tf.name_scope(self.layernorm_before.name):
+ self.layernorm_before.build([None, None, self.dim])
+ if getattr(self, "attention", None) is not None:
+ with tf.name_scope(self.attention.name):
+ self.attention.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_after", None) is not None:
+ with tf.name_scope(self.layernorm_after.name):
+ self.layernorm_after.build([None, None, self.dim])
+ if getattr(self, "intermediate", None) is not None:
+ with tf.name_scope(self.intermediate.name):
+ self.intermediate.build(None)
+ if getattr(self, "swin_output", None) is not None:
+ with tf.name_scope(self.swin_output.name):
+ self.swin_output.build(None)
+
+
+class TFSwinStage(keras.layers.Layer):
+ def __init__(
+ self,
+ config: SwinConfig,
+ dim: int,
+ input_resolution: Tuple[int, int],
+ depth: int,
+ num_heads: int,
+ drop_path: List[float],
+ downsample: Optional[Callable],
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.config = config
+ self.dim = dim
+ self.blocks = [
+ TFSwinLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+ drop_path_rate=drop_path[i],
+ name=f"blocks.{i}",
+ )
+ for i in range(depth)
+ ]
+
+ # patch merging layer
+ if downsample is not None:
+ self.downsample = downsample(
+ input_resolution,
+ dim=dim,
+ norm_layer=partial(keras.layers.LayerNormalization, epsilon=1e-5),
+ name="downsample",
+ )
+ else:
+ self.downsample = None
+
+ self.pointing = False
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ input_dimensions: Tuple[int, int],
+ head_mask: tf.Tensor | None = None,
+ output_attentions: Optional[bool] = False,
+ training: bool = False,
+ ) -> Tuple[tf.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, training=training
+ )
+
+ hidden_states = layer_outputs[0]
+
+ 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(layer_outputs[0], input_dimensions, training=training)
+ else:
+ output_dimensions = (height, width, height, width)
+
+ stage_outputs = (hidden_states, output_dimensions)
+
+ if output_attentions:
+ stage_outputs += layer_outputs[1:]
+ return stage_outputs
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "downsample", None) is not None:
+ with tf.name_scope(self.downsample.name):
+ self.downsample.build(None)
+ if getattr(self, "blocks", None) is not None:
+ for layer in self.blocks:
+ with tf.name_scope(layer.name):
+ layer.build(None)
+
+
+class TFSwinEncoder(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, grid_size: Tuple[int, int], **kwargs):
+ super().__init__(**kwargs)
+ self.num_layers = len(config.depths)
+ self.config = config
+ dpr = list((tf.linspace(0, 1, sum(config.depths)) * config.drop_path_rate).numpy())
+ self.layers = [
+ TFSwinStage(
+ 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=TFSwinPatchMerging if (i_layer < self.num_layers - 1) else None,
+ name=f"layers.{i_layer}",
+ )
+ for i_layer in range(self.num_layers)
+ ]
+
+ self.gradient_checkpointing = False
+
+ def call(
+ self,
+ hidden_states: tf.Tensor,
+ input_dimensions: Tuple[int, int],
+ head_mask: tf.Tensor | None = None,
+ output_attentions: bool = False,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ training: bool = False,
+ ) -> Union[Tuple[tf.Tensor, ...], TFSwinEncoderOutput]:
+ all_input_dimensions = ()
+ 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 = shape_list(hidden_states)
+ # rearrange b (h w) c -> b c h w
+ reshaped_hidden_state = tf.reshape(hidden_states, (batch_size, *input_dimensions, hidden_size))
+ reshaped_hidden_state = tf.transpose(reshaped_hidden_state, (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
+
+ layer_outputs = layer_module(
+ hidden_states, input_dimensions, layer_head_mask, output_attentions, training=training
+ )
+
+ hidden_states = layer_outputs[0]
+ output_dimensions = layer_outputs[1]
+
+ input_dimensions = (output_dimensions[-2], output_dimensions[-1])
+ all_input_dimensions += (input_dimensions,)
+
+ if output_hidden_states:
+ batch_size, _, hidden_size = shape_list(hidden_states)
+ # rearrange b (h w) c -> b c h w
+ reshaped_hidden_state = tf.reshape(hidden_states, (batch_size, *input_dimensions, hidden_size))
+ reshaped_hidden_state = tf.transpose(reshaped_hidden_state, (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 TFSwinEncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ reshaped_hidden_states=all_reshaped_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)
+
+
+class TFSwinPreTrainedModel(TFPreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = SwinConfig
+ base_model_prefix = "swin"
+ main_input_name = "pixel_values"
+
+
+SWIN_START_DOCSTRING = r"""
+ This model is a Tensorflow
+ [keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
+ regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and
+ behavior.
+
+ Parameters:
+ config ([`SwinConfig`]): 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 (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`ViTImageProcessor.__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.
+"""
+
+
+def normalize_data_format(value: str) -> str:
+ """
+ From tensorflow addons
+ https://github.com/tensorflow/addons/blob/8cec33fcaaf1cf90aec7bdd55a0fcdbb251ce5c2/tensorflow_addons/utils/keras_utils.py#L71
+ """
+ if value is None:
+ value = keras.backend.image_data_format()
+ data_format = value.lower()
+ if data_format not in {"channels_first", "channels_last"}:
+ raise ValueError(
+ 'The `data_format` argument must be one of "channels_first", "channels_last". Received: ' + str(value)
+ )
+ return data_format
+
+
+class AdaptiveAveragePooling1D(keras.layers.Layer):
+ """
+ Args:
+ Average 1D Pooling with adaptive kernel size.
+ output_size: An integer or tuple/list of a single integer, specifying pooled_features.
+ The new size of output channels.
+ data_format: A string,
+ one of `channels_last` (default) or `channels_first`. The ordering of the dimensions in the inputs.
+ `channels_last` corresponds to inputs with shape `(batch, steps, channels)` while `channels_first` corresponds
+ to inputs with shape `(batch, channels, steps)`.
+ Input shape:
+ - If `data_format='channels_last'`: 3D tensor with shape `(batch, steps, channels)`.
+ - If `data_format='channels_first'`: 3D tensor with shape `(batch, channels, steps)`.
+ Output shape:
+ - If `data_format='channels_last'`: 3D tensor with shape `(batch_size, pooled_steps, channels)`.
+ - If `data_format='channels_first'`: 3D tensor with shape `(batch_size, channels, pooled_steps)`.
+
+ Adapted from [tensorflow-addon's adaptive pooling.py](
+ https://github.com/tensorflow/addons/blob/8cec33fcaaf1cf90aec7bdd55a0fcdbb251ce5c2/tensorflow_addons/layers/adaptive_pooling.py#L90-L120
+ )
+ """
+
+ def __init__(
+ self,
+ output_size: Union[int, Iterable[int]],
+ reduce_function: Callable = tf.reduce_mean,
+ data_format: Optional[str] = None,
+ **kwargs,
+ ) -> None:
+ self.data_format = normalize_data_format(data_format)
+ self.reduce_function = reduce_function
+ self.output_size = (output_size,) if isinstance(output_size, int) else tuple(output_size)
+ super().__init__(**kwargs)
+
+ def call(self, inputs: tf.Tensor, *args) -> None:
+ bins = self.output_size[0]
+ if self.data_format == "channels_last":
+ splits = tf.split(inputs, bins, axis=1)
+ splits = tf.stack(splits, axis=1)
+ out_vect = self.reduce_function(splits, axis=2)
+ else:
+ splits = tf.split(inputs, bins, axis=2)
+ splits = tf.stack(splits, axis=2)
+ out_vect = self.reduce_function(splits, axis=3)
+ return out_vect
+
+ def compute_output_shape(self, input_shape: Iterable[int]) -> tf.TensorShape:
+ input_shape = tf.TensorShape(input_shape).as_list()
+ if self.data_format == "channels_last":
+ shape = tf.TensorShape([input_shape[0], self.output_size[0], input_shape[2]])
+ else:
+ shape = tf.TensorShape([input_shape[0], input_shape[1], self.output_size[0]])
+ return shape
+
+ def get_config(self) -> Dict[str, Any]:
+ config = {
+ "output_size": self.output_size,
+ "data_format": self.data_format,
+ }
+ base_config = super().get_config()
+ return {**base_config, **config}
+
+
+@keras_serializable
+class TFSwinMainLayer(keras.layers.Layer):
+ config_class = SwinConfig
+
+ def __init__(
+ self, config: SwinConfig, add_pooling_layer: bool = True, use_mask_token: bool = False, **kwargs
+ ) -> None:
+ super().__init__(**kwargs)
+ self.config = config
+ self.num_layers = len(config.depths)
+ self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))
+
+ self.embeddings = TFSwinEmbeddings(config, use_mask_token=use_mask_token, name="embeddings")
+ self.encoder = TFSwinEncoder(config, self.embeddings.patch_grid, name="encoder")
+
+ self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
+ self.pooler = AdaptiveAveragePooling1D(output_size=(1,)) if add_pooling_layer else None
+
+ def get_input_embeddings(self) -> TFSwinPatchEmbeddings:
+ return self.embeddings.patch_embeddings
+
+ def _prune_heads(self, heads_to_prune: Dict[int, List]):
+ """
+ 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 get_head_mask(self, head_mask: Optional[Any]) -> List:
+ if head_mask is not None:
+ raise NotImplementedError
+ return [None] * len(self.config.depths)
+
+ @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[TFSwinModelOutput, 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")
+
+ # 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, input_dimensions = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, training=training
+ )
+
+ 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,
+ training=training,
+ )
+
+ sequence_output = encoder_outputs[0]
+ sequence_output = self.layernorm(sequence_output, training=training)
+
+ pooled_output = None
+ if self.pooler is not None:
+ batch_size, _, num_features = shape_list(sequence_output)
+ pooled_output = self.pooler(sequence_output)
+ pooled_output = tf.reshape(pooled_output, (batch_size, num_features))
+
+ if not return_dict:
+ output = (sequence_output, pooled_output) + encoder_outputs[1:]
+ return output
+
+ return TFSwinModelOutput(
+ 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,
+ )
+
+ 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.num_features])
+
+
+@add_start_docstrings(
+ "The bare Swin Model transformer outputting raw hidden-states without any specific head on top.",
+ SWIN_START_DOCSTRING,
+)
+class TFSwinModel(TFSwinPreTrainedModel):
+ def __init__(
+ self, config: SwinConfig, add_pooling_layer: bool = True, use_mask_token: bool = False, **kwargs
+ ) -> None:
+ super().__init__(config, **kwargs)
+ self.config = config
+ self.swin = TFSwinMainLayer(config, name="swin")
+
+ @add_start_docstrings_to_model_forward(SWIN_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFSwinModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ modality="vision",
+ expected_output=_EXPECTED_OUTPUT_SHAPE,
+ )
+ @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[TFSwinModelOutput, Tuple[tf.Tensor, ...]]:
+ 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).
+ """
+ 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")
+
+ swin_outputs = self.swin(
+ 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 swin_outputs
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "swin", None) is not None:
+ with tf.name_scope(self.swin.name):
+ self.swin.build(None)
+
+
+class TFSwinPixelShuffle(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 TFSwinDecoder(keras.layers.Layer):
+ def __init__(self, config: SwinConfig, **kwargs):
+ super().__init__(**kwargs)
+ self.conv2d = keras.layers.Conv2D(
+ filters=config.encoder_stride**2 * config.num_channels, kernel_size=1, strides=1, name="0"
+ )
+ self.pixel_shuffle = TFSwinPixelShuffle(config.encoder_stride, name="1")
+ self.config = config
+
+ def call(self, x: tf.Tensor) -> tf.Tensor:
+ hidden_states = x
+ # B,C,H,W -> B,H,W,C
+ hidden_states = tf.transpose(hidden_states, (0, 2, 3, 1))
+ hidden_states = self.conv2d(hidden_states)
+ hidden_states = self.pixel_shuffle(hidden_states)
+ # B,H,W,C -> B,C,H,W
+ hidden_states = tf.transpose(hidden_states, (0, 3, 1, 2))
+ 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(
+ "Swin Model with a decoder on top for masked image modeling, as proposed in"
+ " [SimMIM](https://arxiv.org/abs/2111.09886).",
+ SWIN_START_DOCSTRING,
+)
+class TFSwinForMaskedImageModeling(TFSwinPreTrainedModel):
+ def __init__(self, config: SwinConfig):
+ super().__init__(config)
+
+ self.swin = TFSwinMainLayer(config, add_pooling_layer=False, use_mask_token=True, name="swin")
+
+ self.decoder = TFSwinDecoder(config, name="decoder")
+
+ @add_start_docstrings_to_model_forward(SWIN_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFSwinMaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
+ @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, TFSwinMaskedImageModelingOutput]:
+ r"""
+ bool_masked_pos (`tf.Tensor` 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, TFSwinForMaskedImageModeling
+ >>> 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/swin-tiny-patch4-window7-224")
+ >>> model = TFSwinForMaskedImageModeling.from_pretrained("microsoft/swin-tiny-patch4-window7-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.random.uniform((1, num_patches)) >= 0.5
+
+ >>> 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.swin(
+ 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 = tf.transpose(sequence_output, (0, 2, 1))
+ batch_size, num_channels, sequence_length = shape_list(sequence_output)
+ height = width = int(sequence_length**0.5)
+ sequence_output = tf.reshape(sequence_output, (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 = 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[2:]
+ return ((masked_im_loss,) + output) if masked_im_loss is not None else output
+
+ return TFSwinMaskedImageModelingOutput(
+ loss=masked_im_loss,
+ reconstruction=reconstructed_pixel_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "swin", None) is not None:
+ with tf.name_scope(self.swin.name):
+ self.swin.build(None)
+ if getattr(self, "decoder", None) is not None:
+ with tf.name_scope(self.decoder.name):
+ self.decoder.build(None)
+
+
+@add_start_docstrings(
+ """
+ Swin 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.
+ """,
+ SWIN_START_DOCSTRING,
+)
+class TFSwinForImageClassification(TFSwinPreTrainedModel, TFSequenceClassificationLoss):
+ def __init__(self, config: SwinConfig):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.swin = TFSwinMainLayer(config, name="swin")
+
+ # Classifier head
+ self.classifier = (
+ keras.layers.Dense(config.num_labels, name="classifier")
+ if config.num_labels > 0
+ else keras.layers.Activation("linear", name="classifier")
+ )
+
+ @add_start_docstrings_to_model_forward(SWIN_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_IMAGE_CLASS_CHECKPOINT,
+ output_type=TFSwinImageClassifierOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
+ )
+ @unpack_inputs
+ 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[Tuple[tf.Tensor, ...], TFSwinImageClassifierOutput]:
+ 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.swin(
+ 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[1]
+
+ 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[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TFSwinImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "swin", None) is not None:
+ with tf.name_scope(self.swin.name):
+ self.swin.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.swin.num_features])
+
+
+__all__ = ["TFSwinForImageClassification", "TFSwinForMaskedImageModeling", "TFSwinModel", "TFSwinPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/__init__.py b/lib/python3.12/site-packages/transformers/models/swinv2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..172f1573f1eb1bca4893b2495879346c0b044b9c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swinv2/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_swinv2 import *
+ from .modeling_swinv2 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cd830cb3097dea6767678c98f4bb340056aa7cce
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/configuration_swinv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/configuration_swinv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e47f9222777f7cd49adf17be06a8bea9b3214099
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/configuration_swinv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/modeling_swinv2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/modeling_swinv2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4f9a2f0e26614f8bc23e3c083984383ce8e584c0
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/swinv2/__pycache__/modeling_swinv2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/configuration_swinv2.py b/lib/python3.12/site-packages/transformers/models/swinv2/configuration_swinv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..addb30e6a102dcdb835c90fe688045000922f21c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swinv2/configuration_swinv2.py
@@ -0,0 +1,159 @@
+# 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.
+"""Swinv2 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__)
+
+
+class Swinv2Config(BackboneConfigMixin, PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`Swinv2Model`]. It is used to instantiate a Swin
+ Transformer 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 Swin Transformer v2
+ [microsoft/swinv2-tiny-patch4-window8-256](https://huggingface.co/microsoft/swinv2-tiny-patch4-window8-256)
+ 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.
+ pretrained_window_sizes (`list(int)`, *optional*, defaults to `[0, 0, 0, 0]`):
+ Size of windows during pretraining.
+ 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.
+ 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.
+ 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.
+
+ Example:
+
+ ```python
+ >>> from transformers import Swinv2Config, Swinv2Model
+
+ >>> # Initializing a Swinv2 microsoft/swinv2-tiny-patch4-window8-256 style configuration
+ >>> configuration = Swinv2Config()
+
+ >>> # Initializing a model (with random weights) from the microsoft/swinv2-tiny-patch4-window8-256 style configuration
+ >>> model = Swinv2Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "swinv2"
+
+ 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,
+ pretrained_window_sizes=[0, 0, 0, 0],
+ 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,
+ 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.depths = depths
+ self.num_layers = len(depths)
+ self.num_heads = num_heads
+ self.window_size = window_size
+ self.pretrained_window_sizes = pretrained_window_sizes
+ 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
+ self.encoder_stride = encoder_stride
+ 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
+ )
+ # we set the hidden_size attribute in order to make Swinv2 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))
+
+
+__all__ = ["Swinv2Config"]
diff --git a/lib/python3.12/site-packages/transformers/models/swinv2/modeling_swinv2.py b/lib/python3.12/site-packages/transformers/models/swinv2/modeling_swinv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a5d1bd988c13ce4af6048b17ac676e95cca4531
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/swinv2/modeling_swinv2.py
@@ -0,0 +1,1478 @@
+# 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 Swinv2 Transformer model."""
+
+import collections.abc
+import math
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import Tensor, 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, meshgrid, prune_linear_layer
+from ...utils import (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+ torch_int,
+)
+from ...utils.backbone_utils import BackboneMixin
+from .configuration_swinv2 import Swinv2Config
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "Swinv2Config"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "microsoft/swinv2-tiny-patch4-window8-256"
+_EXPECTED_OUTPUT_SHAPE = [1, 64, 768]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "microsoft/swinv2-tiny-patch4-window8-256"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "Egyptian cat"
+
+
+# drop_path, Swinv2PatchEmbeddings, Swinv2PatchMerging and Swinv2DropPath are from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/swin_transformer_v2.py.
+
+
+@dataclass
+# Copied from transformers.models.swin.modeling_swin.SwinEncoderOutput with Swin->Swinv2
+class Swinv2EncoderOutput(ModelOutput):
+ """
+ Swinv2 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: 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.swin.modeling_swin.SwinModelOutput with Swin->Swinv2
+class Swinv2ModelOutput(ModelOutput):
+ """
+ Swinv2 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: Optional[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.swin.modeling_swin.SwinMaskedImageModelingOutput with Swin->Swinv2
+class Swinv2MaskedImageModelingOutput(ModelOutput):
+ """
+ Swinv2 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.
+ 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
+ reconstruction: 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
+
+ @property
+ def logits(self):
+ warnings.warn(
+ "logits attribute is deprecated and will be removed in version 5 of Transformers."
+ " Please use the reconstruction attribute to retrieve the final output instead.",
+ FutureWarning,
+ )
+ return self.reconstruction
+
+
+@dataclass
+# Copied from transformers.models.swin.modeling_swin.SwinImageClassifierOutput with Swin->Swinv2
+class Swinv2ImageClassifierOutput(ModelOutput):
+ """
+ Swinv2 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: 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.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 with Swin->Swinv2
+class Swinv2DropPath(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.SwinEmbeddings with Swin->Swinv2
+class Swinv2Embeddings(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 = Swinv2PatchEmbeddings(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)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding
+ 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. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embeddings.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, :1]
+ patch_pos_embed = self.position_embeddings[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor],
+ bool_masked_pos: Optional[torch.BoolTensor] = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> Tuple[torch.Tensor]:
+ _, num_channels, height, width = pixel_values.shape
+ 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:
+ 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, output_dimensions
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinPatchEmbeddings with Swin->Swinv2
+class Swinv2PatchEmbeddings(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
+ # 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
+
+
+class Swinv2PatchMerging(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(2 * 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.reduction(input_feature)
+ input_feature = self.norm(input_feature)
+
+ return input_feature
+
+
+class Swinv2SelfAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size, pretrained_window_size=[0, 0]):
+ 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.pretrained_window_size = pretrained_window_size
+ self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+ # mlp to generate continuous relative position bias
+ self.continuous_position_bias_mlp = nn.Sequential(
+ nn.Linear(2, 512, bias=True), nn.ReLU(inplace=True), nn.Linear(512, num_heads, bias=False)
+ )
+
+ # get relative_coords_table
+ relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.int64).float()
+ relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.int64).float()
+ relative_coords_table = (
+ torch.stack(meshgrid([relative_coords_h, relative_coords_w], indexing="ij"))
+ .permute(1, 2, 0)
+ .contiguous()
+ .unsqueeze(0)
+ ) # [1, 2*window_height - 1, 2*window_width - 1, 2]
+ if pretrained_window_size[0] > 0:
+ relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
+ relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
+ elif window_size > 1:
+ relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
+ relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
+ relative_coords_table *= 8 # normalize to -8, 8
+ relative_coords_table = (
+ torch.sign(relative_coords_table) * torch.log2(torch.abs(relative_coords_table) + 1.0) / math.log2(8)
+ )
+ # set to same dtype as mlp weight
+ relative_coords_table = relative_coords_table.to(next(self.continuous_position_bias_mlp.parameters()).dtype)
+ self.register_buffer("relative_coords_table", relative_coords_table, persistent=False)
+
+ # 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, persistent=False)
+
+ 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=False)
+ 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)
+
+ # cosine attention
+ attention_scores = nn.functional.normalize(query_layer, dim=-1) @ nn.functional.normalize(
+ key_layer, dim=-1
+ ).transpose(-2, -1)
+ logit_scale = torch.clamp(self.logit_scale, max=math.log(1.0 / 0.01)).exp()
+ attention_scores = attention_scores * logit_scale
+ relative_position_bias_table = self.continuous_position_bias_mlp(self.relative_coords_table).view(
+ -1, self.num_attention_heads
+ )
+ # [window_height*window_width,window_height*window_width,num_attention_heads]
+ relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
+ self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
+ )
+ # [num_attention_heads,window_height*window_width,window_height*window_width]
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+ 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 Swinv2Model 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_mask.unsqueeze(1).unsqueeze(0)
+ 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 with Swin->Swinv2
+class Swinv2SelfOutput(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 Swinv2Attention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size, pretrained_window_size=0):
+ super().__init__()
+ self.self = Swinv2SelfAttention(
+ config=config,
+ dim=dim,
+ num_heads=num_heads,
+ window_size=window_size,
+ pretrained_window_size=pretrained_window_size
+ if isinstance(pretrained_window_size, collections.abc.Iterable)
+ else (pretrained_window_size, pretrained_window_size),
+ )
+ self.output = Swinv2SelfOutput(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 with Swin->Swinv2
+class Swinv2Intermediate(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 with Swin->Swinv2
+class Swinv2Output(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 Swinv2Layer(nn.Module):
+ def __init__(
+ self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0, pretrained_window_size=0
+ ):
+ super().__init__()
+ self.input_resolution = input_resolution
+ window_size, shift_size = self._compute_window_shift(
+ (config.window_size, config.window_size), (shift_size, shift_size)
+ )
+ self.window_size = window_size[0]
+ self.shift_size = shift_size[0]
+ self.attention = Swinv2Attention(
+ config=config,
+ dim=dim,
+ num_heads=num_heads,
+ window_size=self.window_size,
+ pretrained_window_size=pretrained_window_size
+ if isinstance(pretrained_window_size, collections.abc.Iterable)
+ else (pretrained_window_size, pretrained_window_size),
+ )
+ self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.drop_path = Swinv2DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+ self.intermediate = Swinv2Intermediate(config, dim)
+ self.output = Swinv2Output(config, dim)
+ self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+
+ def _compute_window_shift(self, target_window_size, target_shift_size) -> Tuple[Tuple[int, int], Tuple[int, int]]:
+ window_size = [r if r <= w else w for r, w in zip(self.input_resolution, target_window_size)]
+ shift_size = [0 if r <= w else s for r, w, s in zip(self.input_resolution, window_size, target_shift_size)]
+ return window_size, shift_size
+
+ def get_attn_mask(self, height, width, dtype):
+ if self.shift_size > 0:
+ # calculate attention mask for shifted window multihead self attention
+ 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,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ height, width = input_dimensions
+ batch_size, _, channels = hidden_states.size()
+ shortcut = hidden_states
+
+ # pad hidden_states to multiples of window size
+ hidden_states = hidden_states.view(batch_size, height, width, channels)
+ 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 = self.layernorm_before(attention_windows)
+ hidden_states = shortcut + self.drop_path(hidden_states)
+
+ layer_output = self.intermediate(hidden_states)
+ layer_output = self.output(layer_output)
+ layer_output = hidden_states + self.drop_path(self.layernorm_after(layer_output))
+
+ layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
+ return layer_outputs
+
+
+class Swinv2Stage(nn.Module):
+ def __init__(
+ self, config, dim, input_resolution, depth, num_heads, drop_path, downsample, pretrained_window_size=0
+ ):
+ super().__init__()
+ self.config = config
+ self.dim = dim
+ blocks = []
+ for i in range(depth):
+ block = Swinv2Layer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+ pretrained_window_size=pretrained_window_size,
+ )
+ blocks.append(block)
+ self.blocks = nn.ModuleList(blocks)
+
+ # 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,
+ ) -> 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,
+ )
+
+ 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
+
+
+class Swinv2Encoder(nn.Module):
+ def __init__(self, config, grid_size, pretrained_window_sizes=(0, 0, 0, 0)):
+ super().__init__()
+ self.num_layers = len(config.depths)
+ self.config = config
+ if self.config.pretrained_window_sizes is not None:
+ pretrained_window_sizes = config.pretrained_window_sizes
+ dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
+
+ layers = []
+ for i_layer in range(self.num_layers):
+ stage = Swinv2Stage(
+ 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=Swinv2PatchMerging if (i_layer < self.num_layers - 1) else None,
+ pretrained_window_size=pretrained_window_sizes[i_layer],
+ )
+ layers.append(stage)
+ self.layers = nn.ModuleList(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,
+ return_dict: Optional[bool] = True,
+ ) -> Union[Tuple, Swinv2EncoderOutput]:
+ 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
+ )
+ else:
+ layer_outputs = layer_module(
+ hidden_states,
+ input_dimensions,
+ layer_head_mask,
+ output_attentions,
+ )
+
+ 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, all_reshaped_hidden_states]
+ if v is not None
+ )
+
+ return Swinv2EncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ reshaped_hidden_states=all_reshaped_hidden_states,
+ )
+
+
+class Swinv2PreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = Swinv2Config
+ base_model_prefix = "swinv2"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Swinv2Stage"]
+
+ 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)
+ elif isinstance(module, Swinv2Embeddings):
+ if module.mask_token is not None:
+ module.mask_token.data.zero_()
+ if module.position_embeddings is not None:
+ module.position_embeddings.data.zero_()
+ elif isinstance(module, Swinv2SelfAttention):
+ module.logit_scale.data.fill_(math.log(10))
+
+
+SWINV2_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 ([`Swinv2Config`]): 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.
+"""
+
+SWINV2_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.
+ 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.
+ interpolate_pos_encoding (`bool`, *optional*, default `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare Swinv2 Model transformer outputting raw hidden-states without any specific head on top.",
+ SWINV2_START_DOCSTRING,
+)
+# Copied from transformers.models.swin.modeling_swin.SwinModel with SWIN->SWINV2,Swin->Swinv2
+class Swinv2Model(Swinv2PreTrainedModel):
+ 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 = Swinv2Embeddings(config, use_mask_token=use_mask_token)
+ self.encoder = Swinv2Encoder(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
+
+ 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(SWINV2_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=Swinv2ModelOutput,
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, Swinv2ModelOutput]:
+ 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, len(self.config.depths))
+
+ embedding_output, input_dimensions = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ 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]
+ 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 Swinv2ModelOutput(
+ 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(
+ """Swinv2 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).
+
+
+ """,
+ SWINV2_START_DOCSTRING,
+)
+# Copied from transformers.models.swin.modeling_swin.SwinForMaskedImageModeling with swin->swinv2, base-simmim-window6-192->tiny-patch4-window8-256,SWIN->SWINV2,Swin->Swinv2,192->256
+class Swinv2ForMaskedImageModeling(Swinv2PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.swinv2 = Swinv2Model(config, add_pooling_layer=False, use_mask_token=True)
+
+ num_features = int(config.embed_dim * 2 ** (config.num_layers - 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(SWINV2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=Swinv2MaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, Swinv2MaskedImageModelingOutput]:
+ 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, Swinv2ForMaskedImageModeling
+ >>> 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/swinv2-tiny-patch4-window8-256")
+ >>> model = Swinv2ForMaskedImageModeling.from_pretrained("microsoft/swinv2-tiny-patch4-window8-256")
+
+ >>> 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, 256, 256]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.swinv2(
+ pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ 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 Swinv2MaskedImageModelingOutput(
+ loss=masked_im_loss,
+ reconstruction=reconstructed_pixel_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@add_start_docstrings(
+ """
+ Swinv2 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.
+
+
+
+ Note that it's possible to fine-tune SwinV2 on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """,
+ SWINV2_START_DOCSTRING,
+)
+# Copied from transformers.models.swin.modeling_swin.SwinForImageClassification with SWIN->SWINV2,Swin->Swinv2,swin->swinv2
+class Swinv2ForImageClassification(Swinv2PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.swinv2 = Swinv2Model(config)
+
+ # Classifier head
+ self.classifier = (
+ nn.Linear(self.swinv2.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(SWINV2_INPUTS_DOCSTRING)
+ @add_code_sample_docstrings(
+ checkpoint=_IMAGE_CLASS_CHECKPOINT,
+ output_type=Swinv2ImageClassifierOutput,
+ config_class=_CONFIG_FOR_DOC,
+ expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
+ )
+ 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, Swinv2ImageClassifierOutput]:
+ 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.swinv2(
+ pixel_values,
+ 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 = 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 Swinv2ImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@add_start_docstrings(
+ """
+ Swinv2 backbone, to be used with frameworks like DETR and MaskFormer.
+ """,
+ SWINV2_START_DOCSTRING,
+)
+class Swinv2Backbone(Swinv2PreTrainedModel, BackboneMixin):
+ def __init__(self, config):
+ super().__init__(config)
+ super()._init_backbone(config)
+
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+ self.embeddings = Swinv2Embeddings(config)
+ self.encoder = Swinv2Encoder(config, self.embeddings.patch_grid)
+
+ # 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(SWINV2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ pixel_values: Tensor,
+ output_attentions: Optional[bool] = None,
+ 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/swinv2-tiny-patch4-window8-256")
+ >>> model = AutoBackbone.from_pretrained(
+ ... "microsoft/swinv2-tiny-patch4-window8-256", 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, 2048, 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, input_dimensions = self.embeddings(pixel_values)
+
+ outputs = self.encoder(
+ embedding_output,
+ input_dimensions,
+ head_mask=None,
+ output_attentions=output_attentions,
+ output_hidden_states=True,
+ output_hidden_states_before_downsampling=True,
+ return_dict=return_dict,
+ )
+
+ hidden_states = outputs.reshaped_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:
+ feature_maps += (hidden_state,)
+
+ if not return_dict:
+ output = (feature_maps,)
+ if output_hidden_states:
+ output += (outputs[1],)
+ if output_attentions:
+ output += (outputs[2],)
+ return output
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.hidden_states if output_hidden_states else None,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "Swinv2ForImageClassification",
+ "Swinv2ForMaskedImageModeling",
+ "Swinv2Model",
+ "Swinv2PreTrainedModel",
+ "Swinv2Backbone",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/__init__.py b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fbc4150412a73ec6d77a09db13c105e95e50c06
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__init__.py
@@ -0,0 +1,28 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_timm_wrapper import *
+ from .modeling_timm_wrapper import *
+ from .processing_timm_wrapper import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e62ae66b972478fd3288f20faa9d79eefebecc25
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/configuration_timm_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/configuration_timm_wrapper.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..795ac12505f7eb7211608c04c05b205a15474556
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/configuration_timm_wrapper.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/image_processing_timm_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/image_processing_timm_wrapper.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2befc2ba950d8c3a3d5844c3ce0f26de04d848c8
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/image_processing_timm_wrapper.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/modeling_timm_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/modeling_timm_wrapper.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c80ed6d132e15f55a64bdd516e7b3d6d8c36cb15
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/timm_wrapper/__pycache__/modeling_timm_wrapper.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/configuration_timm_wrapper.py b/lib/python3.12/site-packages/transformers/models/timm_wrapper/configuration_timm_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..4313ad854add5955150268d43d4f3a6b97e73ed1
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/timm_wrapper/configuration_timm_wrapper.py
@@ -0,0 +1,117 @@
+# 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.
+
+"""Configuration for TimmWrapper models"""
+
+from typing import Any, Dict
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import is_timm_available, logging, requires_backends
+
+
+if is_timm_available():
+ from timm.data import ImageNetInfo, infer_imagenet_subset
+
+
+logger = logging.get_logger(__name__)
+
+
+class TimmWrapperConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration for a timm backbone [`TimmWrapper`].
+
+ It is used to instantiate a timm model according to the specified arguments, defining the model.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Config loads imagenet label descriptions and stores them in `id2label` attribute, `label2id` attribute for default
+ imagenet models is set to `None` due to occlusions in the label descriptions.
+
+ Args:
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ do_pooling (`bool`, *optional*, defaults to `True`):
+ Whether to do pooling for the last_hidden_state in `TimmWrapperModel` or not.
+
+ Example:
+ ```python
+ >>> from transformers import TimmWrapperModel
+
+ >>> # Initializing a timm model
+ >>> model = TimmWrapperModel.from_pretrained("timm/resnet18.a1_in1k")
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "timm_wrapper"
+
+ def __init__(self, initializer_range: float = 0.02, do_pooling: bool = True, **kwargs):
+ self.initializer_range = initializer_range
+ self.do_pooling = do_pooling
+ super().__init__(**kwargs)
+
+ @classmethod
+ def from_dict(cls, config_dict: Dict[str, Any], **kwargs):
+ label_names = config_dict.get("label_names", None)
+ is_custom_model = "num_labels" in kwargs or "id2label" in kwargs
+
+ # if no labels added to config, use imagenet labeller in timm
+ if label_names is None and not is_custom_model:
+ requires_backends(cls, ["timm"])
+ imagenet_subset = infer_imagenet_subset(config_dict)
+ if imagenet_subset:
+ dataset_info = ImageNetInfo(imagenet_subset)
+ synsets = dataset_info.label_names()
+ label_descriptions = dataset_info.label_descriptions(as_dict=True)
+ label_names = [label_descriptions[synset] for synset in synsets]
+
+ if label_names is not None and not is_custom_model:
+ kwargs["id2label"] = dict(enumerate(label_names))
+
+ # if all label names are unique, create label2id mapping as well
+ if len(set(label_names)) == len(label_names):
+ kwargs["label2id"] = {name: i for i, name in enumerate(label_names)}
+ else:
+ kwargs["label2id"] = None
+
+ # timm config stores the `num_classes` attribute in both the root of config and in the "pretrained_cfg" dict.
+ # We are removing these attributes in order to have the native `transformers` num_labels attribute in config
+ # and to avoid duplicate attributes
+ num_labels_in_kwargs = kwargs.pop("num_labels", None)
+ num_labels_in_dict = config_dict.pop("num_classes", None)
+
+ # passed num_labels has priority over num_classes in config_dict
+ kwargs["num_labels"] = num_labels_in_kwargs or num_labels_in_dict
+
+ # pop num_classes from "pretrained_cfg",
+ # it is not necessary to have it, only root one is used in timm
+ if "pretrained_cfg" in config_dict and "num_classes" in config_dict["pretrained_cfg"]:
+ config_dict["pretrained_cfg"].pop("num_classes", None)
+
+ return super().from_dict(config_dict, **kwargs)
+
+ def to_dict(self) -> Dict[str, Any]:
+ output = super().to_dict()
+ output["num_classes"] = self.num_labels
+ output["label_names"] = list(self.id2label.values())
+ output.pop("id2label", None)
+ output.pop("label2id", None)
+ return output
+
+
+__all__ = ["TimmWrapperConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/image_processing_timm_wrapper.py b/lib/python3.12/site-packages/transformers/models/timm_wrapper/image_processing_timm_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..02075a50fb26763c5ca4698e0598e77692a35a51
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/timm_wrapper/image_processing_timm_wrapper.py
@@ -0,0 +1,138 @@
+# 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.
+
+import os
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature
+from ...image_transforms import to_pil_image
+from ...image_utils import ImageInput, make_list_of_images
+from ...utils import TensorType, logging, requires_backends
+from ...utils.import_utils import is_timm_available, is_torch_available
+
+
+if is_timm_available():
+ import timm
+
+if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class TimmWrapperImageProcessor(BaseImageProcessor):
+ """
+ Wrapper class for timm models to be used within transformers.
+
+ Args:
+ pretrained_cfg (`Dict[str, Any]`):
+ The configuration of the pretrained model used to resolve evaluation and
+ training transforms.
+ architecture (`Optional[str]`, *optional*):
+ Name of the architecture of the model.
+ """
+
+ main_input_name = "pixel_values"
+
+ def __init__(
+ self,
+ pretrained_cfg: Dict[str, Any],
+ architecture: Optional[str] = None,
+ **kwargs,
+ ):
+ requires_backends(self, "timm")
+ super().__init__(architecture=architecture)
+
+ self.data_config = timm.data.resolve_data_config(pretrained_cfg, model=None, verbose=False)
+ self.val_transforms = timm.data.create_transform(**self.data_config, is_training=False)
+
+ # useful for training, see examples/pytorch/image-classification/run_image_classification.py
+ self.train_transforms = timm.data.create_transform(**self.data_config, is_training=True)
+
+ # If `ToTensor` is in the transforms, then the input should be numpy array or PIL image.
+ # Otherwise, the input can be a tensor. In later timm versions, `MaybeToTensor` is used
+ # which can handle both numpy arrays / PIL images and tensors.
+ self._not_supports_tensor_input = any(
+ transform.__class__.__name__ == "ToTensor" for transform in self.val_transforms.transforms
+ )
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary.
+ """
+ output = super().to_dict()
+ output.pop("train_transforms", None)
+ output.pop("val_transforms", None)
+ output.pop("_not_supports_tensor_input", None)
+ return output
+
+ @classmethod
+ def get_image_processor_dict(
+ cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+ ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ """
+ Get the image processor dict for the model.
+ """
+ image_processor_filename = kwargs.pop("image_processor_filename", "config.json")
+ return super().get_image_processor_dict(
+ pretrained_model_name_or_path, image_processor_filename=image_processor_filename, **kwargs
+ )
+
+ def preprocess(
+ self,
+ images: ImageInput,
+ return_tensors: Optional[Union[str, TensorType]] = "pt",
+ ) -> BatchFeature:
+ """
+ Preprocess an image or batch of images.
+
+ Args:
+ images (`ImageInput`):
+ Image to preprocess. Expects a single or batch of images
+ return_tensors (`str` or `TensorType`, *optional*):
+ The type of tensors to return.
+ """
+ if return_tensors != "pt":
+ raise ValueError(f"return_tensors for TimmWrapperImageProcessor must be 'pt', but got {return_tensors}")
+
+ if self._not_supports_tensor_input and isinstance(images, torch.Tensor):
+ images = images.cpu().numpy()
+
+ # If the input is a torch tensor, then no conversion is needed
+ # Otherwise, we need to pass in a list of PIL images
+ if isinstance(images, torch.Tensor):
+ images = self.val_transforms(images)
+ # Add batch dimension if a single image
+ images = images.unsqueeze(0) if images.ndim == 3 else images
+ else:
+ images = make_list_of_images(images)
+ images = [to_pil_image(image) for image in images]
+ images = torch.stack([self.val_transforms(image) for image in images])
+
+ return BatchFeature({"pixel_values": images}, tensor_type=return_tensors)
+
+ def save_pretrained(self, *args, **kwargs):
+ # disable it to make checkpoint the same as in `timm` library.
+ logger.warning_once(
+ "The `save_pretrained` method is disabled for TimmWrapperImageProcessor. "
+ "The image processor configuration is saved directly in `config.json` when "
+ "`save_pretrained` is called for saving the model."
+ )
+
+
+__all__ = ["TimmWrapperImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/timm_wrapper/modeling_timm_wrapper.py b/lib/python3.12/site-packages/transformers/models/timm_wrapper/modeling_timm_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..48c1d85825335ba66f4328d1db9fa5e191f46ecc
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/timm_wrapper/modeling_timm_wrapper.py
@@ -0,0 +1,367 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...modeling_outputs import ImageClassifierOutput, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_start_docstrings_to_model_forward,
+ is_timm_available,
+ replace_return_docstrings,
+ requires_backends,
+)
+from .configuration_timm_wrapper import TimmWrapperConfig
+
+
+if is_timm_available():
+ import timm
+
+
+@dataclass
+class TimmWrapperModelOutput(ModelOutput):
+ """
+ Output class for models TimmWrapperModel, containing the last hidden states, an optional pooled output,
+ and optional hidden states.
+
+ Args:
+ last_hidden_state (`torch.FloatTensor`):
+ The last hidden state of the model, output before applying the classification head.
+ pooler_output (`torch.FloatTensor`, *optional*):
+ The pooled output derived from the last hidden state, if applicable.
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+ A tuple containing the intermediate hidden states of the model at the output of each layer or specified layers.
+ Returned if `output_hidden_states=True` is set or if `config.output_hidden_states=True`.
+ attentions (`tuple(torch.FloatTensor)`, *optional*):
+ A tuple containing the intermediate attention weights of the model at the output of each layer.
+ Returned if `output_attentions=True` is set or if `config.output_attentions=True`.
+ Note: Currently, Timm models do not support attentions output.
+ """
+
+ last_hidden_state: torch.FloatTensor
+ pooler_output: Optional[torch.FloatTensor] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+TIMM_WRAPPER_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 [`TimmWrapperImageProcessor.preprocess`]
+ for details.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. Not compatible with timm wrapped models.
+ output_hidden_states (`bool`, *optional*):
+ Whether or not to return the hidden states of all layers. Not compatible with timm wrapped models.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+ **kwargs:
+ Additional keyword arguments passed along to the `timm` model forward.
+"""
+
+
+class TimmWrapperPreTrainedModel(PreTrainedModel):
+ main_input_name = "pixel_values"
+ config_class = TimmWrapperConfig
+ _no_split_modules = []
+ model_tags = ["timm"]
+
+ # used in Trainer to avoid passing `loss_kwargs` to model forward
+ accepts_loss_kwargs = False
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["vision", "timm"])
+ super().__init__(*args, **kwargs)
+
+ @staticmethod
+ def _fix_state_dict_key_on_load(key) -> Tuple[str, bool]:
+ """
+ Overrides original method that renames `gamma` and `beta` to `weight` and `bias`.
+ We don't want this behavior for timm wrapped models. Instead, this method adds a
+ "timm_model." prefix to enable loading official timm Hub checkpoints.
+ """
+ if "timm_model." not in key:
+ return f"timm_model.{key}", True
+ return key, False
+
+ def _fix_state_dict_key_on_save(self, key):
+ """
+ Overrides original method to remove "timm_model." prefix from state_dict keys.
+ Makes the saved checkpoint compatible with the `timm` library.
+ """
+ return key.replace("timm_model.", ""), True
+
+ def load_state_dict(self, state_dict, *args, **kwargs):
+ """
+ Override original method to fix state_dict keys on load for cases when weights are loaded
+ without using the `from_pretrained` method (e.g., in Trainer to resume from checkpoint).
+ """
+ state_dict = {self._fix_state_dict_key_on_load(k)[0]: v for k, v in state_dict.items()}
+ return super().load_state_dict(state_dict, *args, **kwargs)
+
+ def _init_weights(self, module):
+ """
+ Initialize weights function to properly initialize Linear layer weights.
+ Since model architectures may vary, we assume only the classifier requires
+ initialization, while all other weights should be loaded from the checkpoint.
+ """
+ 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_()
+
+
+class TimmWrapperModel(TimmWrapperPreTrainedModel):
+ """
+ Wrapper class for timm models to be used in transformers.
+ """
+
+ def __init__(self, config: TimmWrapperConfig):
+ super().__init__(config)
+ # using num_classes=0 to avoid creating classification head
+ self.timm_model = timm.create_model(config.architecture, pretrained=False, num_classes=0)
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(TIMM_WRAPPER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TimmWrapperModelOutput, config_class=TimmWrapperConfig)
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[Union[bool, List[int]]] = None,
+ return_dict: Optional[bool] = None,
+ do_pooling: Optional[bool] = None,
+ **kwargs,
+ ) -> Union[TimmWrapperModelOutput, Tuple[Tensor, ...]]:
+ r"""
+ do_pooling (`bool`, *optional*):
+ Whether to do pooling for the last_hidden_state in `TimmWrapperModel` or not. If `None` is passed, the
+ `do_pooling` value from the config is used.
+
+ Returns:
+
+ Examples:
+ ```python
+ >>> import torch
+ >>> from PIL import Image
+ >>> from urllib.request import urlopen
+ >>> from transformers import AutoModel, AutoImageProcessor
+
+ >>> # Load image
+ >>> image = Image.open(urlopen(
+ ... 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
+ ... ))
+
+ >>> # Load model and image processor
+ >>> checkpoint = "timm/resnet50.a1_in1k"
+ >>> image_processor = AutoImageProcessor.from_pretrained(checkpoint)
+ >>> model = AutoModel.from_pretrained(checkpoint).eval()
+
+ >>> # Preprocess image
+ >>> inputs = image_processor(image)
+
+ >>> # Forward pass
+ >>> with torch.no_grad():
+ ... outputs = model(**inputs)
+
+ >>> # Get pooled output
+ >>> pooled_output = outputs.pooler_output
+
+ >>> # Get last hidden state
+ >>> last_hidden_state = outputs.last_hidden_state
+ ```
+ """
+ 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
+ do_pooling = do_pooling if do_pooling is not None else self.config.do_pooling
+
+ if output_attentions:
+ raise ValueError("Cannot set `output_attentions` for timm models.")
+
+ if output_hidden_states and not hasattr(self.timm_model, "forward_intermediates"):
+ raise ValueError(
+ "The 'output_hidden_states' option cannot be set for this timm model. "
+ "To enable this feature, the 'forward_intermediates' method must be implemented "
+ "in the timm model (available in timm versions > 1.*). Please consider using a "
+ "different architecture or updating the timm package to a compatible version."
+ )
+
+ pixel_values = pixel_values.to(self.device, self.dtype)
+
+ if output_hidden_states:
+ # to enable hidden states selection
+ if isinstance(output_hidden_states, (list, tuple)):
+ kwargs["indices"] = output_hidden_states
+ last_hidden_state, hidden_states = self.timm_model.forward_intermediates(pixel_values, **kwargs)
+ else:
+ last_hidden_state = self.timm_model.forward_features(pixel_values, **kwargs)
+ hidden_states = None
+
+ if do_pooling:
+ # classification head is not created, applying pooling only
+ pooler_output = self.timm_model.forward_head(last_hidden_state)
+ else:
+ pooler_output = None
+
+ if not return_dict:
+ outputs = (last_hidden_state, pooler_output, hidden_states)
+ outputs = tuple(output for output in outputs if output is not None)
+ return outputs
+
+ return TimmWrapperModelOutput(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooler_output,
+ hidden_states=hidden_states,
+ )
+
+
+class TimmWrapperForImageClassification(TimmWrapperPreTrainedModel):
+ """
+ Wrapper class for timm models to be used in transformers for image classification.
+ """
+
+ def __init__(self, config: TimmWrapperConfig):
+ super().__init__(config)
+
+ if config.num_labels == 0:
+ raise ValueError(
+ "You are trying to load weights into `TimmWrapperForImageClassification` from a checkpoint with no classifier head. "
+ "Please specify the number of classes, e.g. `model = TimmWrapperForImageClassification.from_pretrained(..., num_labels=10)`, "
+ "or use `TimmWrapperModel` for feature extraction."
+ )
+
+ self.timm_model = timm.create_model(config.architecture, pretrained=False, num_classes=config.num_labels)
+ self.num_labels = config.num_labels
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(TIMM_WRAPPER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=ImageClassifierOutput, config_class=TimmWrapperConfig)
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ labels: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[Union[bool, List[int]]] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs,
+ ) -> Union[ImageClassifierOutput, Tuple[Tensor, ...]]:
+ 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
+ >>> import torch
+ >>> from PIL import Image
+ >>> from urllib.request import urlopen
+ >>> from transformers import AutoModelForImageClassification, AutoImageProcessor
+
+ >>> # Load image
+ >>> image = Image.open(urlopen(
+ ... 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
+ ... ))
+
+ >>> # Load model and image processor
+ >>> checkpoint = "timm/resnet50.a1_in1k"
+ >>> image_processor = AutoImageProcessor.from_pretrained(checkpoint)
+ >>> model = AutoModelForImageClassification.from_pretrained(checkpoint).eval()
+
+ >>> # Preprocess image
+ >>> inputs = image_processor(image)
+
+ >>> # Forward pass
+ >>> with torch.no_grad():
+ ... logits = model(**inputs).logits
+
+ >>> # Get top 5 predictions
+ >>> top5_probabilities, top5_class_indices = torch.topk(logits.softmax(dim=1) * 100, k=5)
+ ```
+ """
+ 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 output_attentions:
+ raise ValueError("Cannot set `output_attentions` for timm models.")
+
+ if output_hidden_states and not hasattr(self.timm_model, "forward_intermediates"):
+ raise ValueError(
+ "The 'output_hidden_states' option cannot be set for this timm model. "
+ "To enable this feature, the 'forward_intermediates' method must be implemented "
+ "in the timm model (available in timm versions > 1.*). Please consider using a "
+ "different architecture or updating the timm package to a compatible version."
+ )
+
+ pixel_values = pixel_values.to(self.device, self.dtype)
+
+ if output_hidden_states:
+ # to enable hidden states selection
+ if isinstance(output_hidden_states, (list, tuple)):
+ kwargs["indices"] = output_hidden_states
+ last_hidden_state, hidden_states = self.timm_model.forward_intermediates(pixel_values, **kwargs)
+ logits = self.timm_model.forward_head(last_hidden_state)
+ else:
+ logits = self.timm_model(pixel_values, **kwargs)
+ hidden_states = None
+
+ 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:
+ outputs = (loss, logits, hidden_states)
+ outputs = tuple(output for output in outputs if output is not None)
+ return outputs
+
+ return ImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=hidden_states,
+ )
+
+
+__all__ = ["TimmWrapperPreTrainedModel", "TimmWrapperModel", "TimmWrapperForImageClassification"]
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/__init__.py b/lib/python3.12/site-packages/transformers/models/vipllava/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ff5c6838dfaff4eb96663c6e2a808e5f082f0de
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vipllava/__init__.py
@@ -0,0 +1,27 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_vipllava import *
+ from .modeling_vipllava import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c771b7c310e46f103832626bae55539b68fc822d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/configuration_vipllava.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/configuration_vipllava.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e5549c876e4763e46756a3f0beadc8c60213f784
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/configuration_vipllava.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/modeling_vipllava.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/modeling_vipllava.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..95830a89cfbdab98ab0e2e2975ffeabbcc7b58da
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vipllava/__pycache__/modeling_vipllava.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/configuration_vipllava.py b/lib/python3.12/site-packages/transformers/models/vipllava/configuration_vipllava.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac24cce24129eb783359a345fc40bdd3016e1d0c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vipllava/configuration_vipllava.py
@@ -0,0 +1,121 @@
+# coding=utf-8
+# Copyright 2023 Microsoft Research & University of Wisconsin-Madison 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.
+"""VipLlava model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class VipLlavaConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`VipLlavaForConditionalGeneration`]. It is used to instantiate an
+ VipLlava 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 VipLlava-9B.
+
+ e.g. [ybelkada/vip-llava-7b-hf](https://huggingface.co/ybelkada/vip-llava-7b-hf)
+
+ 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 (`VipLlavaVisionConfig`, *optional*):
+ Custom vision config or dict
+ text_config (`Union[AutoConfig, dict]`, *optional*):
+ The config object of the text backbone. Can be any of `LlamaConfig` or `MistralConfig`.
+ image_token_index (`int`, *optional*, defaults to 32000):
+ The image token index to encode the image prompt.
+ projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+ The activation function used by the multimodal projector.
+ projector_layernorm_eps (`float`, *optional*, defaults to 1e-05):
+ The layer norm epsilon of the projector layernorm
+ vision_feature_layers (`Union[int, List[int]]`, *optional*, defaults to `[-2, -5, -8, -11, 6]`):
+ The vision feature layer, or list of layers to select the vision features from.
+ image_seq_length (`int`, *optional*, defaults to 576):
+ Sequence length of one image embedding.
+
+ Example:
+
+ ```python
+ >>> from transformers import VipLlavaForConditionalGeneration, VipLlavaConfig, CLIPVisionConfig, LlamaConfig
+
+ >>> # Initializing a CLIP-vision config
+ >>> vision_config = CLIPVisionConfig()
+
+ >>> # Initializing a Llama config
+ >>> text_config = LlamaConfig()
+
+ >>> # Initializing a VipLlava vipllava-7b style configuration
+ >>> configuration = VipLlavaConfig(vision_config, text_config)
+
+ >>> # Initializing a model from the vipllava-7b style configuration
+ >>> model = VipLlavaForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "vipllava"
+ sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
+
+ def __init__(
+ self,
+ vision_config=None,
+ text_config=None,
+ image_token_index=32000,
+ projector_hidden_act="gelu",
+ projector_layernorm_eps=1e-5,
+ vision_feature_layers=[-2, -5, -8, -11, 6],
+ image_seq_length=576,
+ **kwargs,
+ ):
+ self.image_token_index = image_token_index
+ self.projector_hidden_act = projector_hidden_act
+ self.projector_layernorm_eps = projector_layernorm_eps
+ self.vision_feature_layers = vision_feature_layers
+ self.image_seq_length = image_seq_length
+ self.vision_config = vision_config
+
+ if isinstance(self.vision_config, dict):
+ vision_config["model_type"] = (
+ vision_config["model_type"] if "model_type" in vision_config else "clip_vision_model"
+ )
+ self.vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
+ elif vision_config is None:
+ self.vision_config = CONFIG_MAPPING["clip_vision_model"](
+ intermediate_size=4096,
+ hidden_size=1024,
+ patch_size=14,
+ image_size=336,
+ num_hidden_layers=24,
+ num_attention_heads=16,
+ vocab_size=32000,
+ projection_dim=768,
+ )
+
+ if isinstance(text_config, dict):
+ text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama"
+ text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+ elif text_config is None:
+ text_config = CONFIG_MAPPING["llama"]()
+
+ self.text_config = text_config
+
+ super().__init__(**kwargs)
+
+
+__all__ = ["VipLlavaConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/vipllava/modeling_vipllava.py b/lib/python3.12/site-packages/transformers/models/vipllava/modeling_vipllava.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c706e43d5994c7a58ad427194e29ba4b0940131
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vipllava/modeling_vipllava.py
@@ -0,0 +1,478 @@
+# 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.
+"""PyTorch VipLlava model."""
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_outputs import ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_torchdynamo_compiling,
+ logging,
+ replace_return_docstrings,
+)
+from ...utils.deprecation import deprecate_kwarg
+from ..auto import AutoModel, AutoModelForCausalLM
+from .configuration_vipllava import VipLlavaConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "VipLlavaConfig"
+
+
+@dataclass
+# Copied from transformers.models.llava.modeling_llava.LlavaCausalLMOutputWithPast with Llava->VipLlava
+class VipLlavaCausalLMOutputWithPast(ModelOutput):
+ """
+ Base class for VipLlava causal language model (or autoregressive) outputs.
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token 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).
+ 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)`)
+
+ 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.
+ 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_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size (batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
+ image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+class VipLlavaMultiModalProjector(nn.Module):
+ def __init__(self, config: VipLlavaConfig):
+ super().__init__()
+ num_feature_layers = 1 if isinstance(config.vision_feature_layers, int) else len(config.vision_feature_layers)
+ self.projector_layernorm = nn.LayerNorm(
+ num_feature_layers * config.vision_config.hidden_size, eps=config.projector_layernorm_eps
+ )
+
+ self.linear_1 = nn.Linear(
+ num_feature_layers * config.vision_config.hidden_size,
+ config.text_config.hidden_size,
+ bias=True,
+ )
+ self.act = ACT2FN[config.projector_hidden_act]
+ self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+
+ def forward(self, hidden_states):
+ hidden_states = self.projector_layernorm(hidden_states)
+ hidden_states = self.linear_1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+
+VIPLLAVA_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 ([`VipLlavaConfig`] or [`VipLlavaVisionConfig`]):
+ 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 VipLlava Model outputting raw hidden-states without any specific head on top.",
+ VIPLLAVA_START_DOCSTRING,
+)
+# Copied from transformers.models.llava.modeling_llava.LlavaPreTrainedModel with Llava->VipLlava,llava->vipllava
+class VipLlavaPreTrainedModel(PreTrainedModel):
+ config_class = VipLlavaConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["VipLlavaVisionAttention"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_cache_class = True
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+ _supports_quantized_cache = True
+ _supports_static_cache = True
+
+ def _init_weights(self, module):
+ # important: this ported version of VipLlava isn't meant for training from scratch - only
+ # inference and fine-tuning - so the proper init weights code has been removed - the original codebase
+ # https://github.com/haotian-liu/LLaVA/tree/main/vipllava should serve for that purpose
+ std = (
+ self.config.initializer_range
+ if hasattr(self.config, "initializer_range")
+ else self.config.text_config.initializer_range
+ )
+
+ if hasattr(module, "class_embedding"):
+ module.class_embedding.data.normal_(mean=0.0, std=std)
+
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ 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_()
+
+
+VIPLLAVA_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)
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
+ The tensors corresponding to the input images. Pixel values can be obtained using
+ [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details ([]`LlavaProcessor`] uses
+ [`CLIPImageProcessor`] for processing images).
+ 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.
+ 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 VIPLLAVA model which consists of a vision backbone and a language model.""",
+ VIPLLAVA_START_DOCSTRING,
+)
+# Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration with LLAVA->VIPLLAVA,Llava->VipLlava
+class VipLlavaForConditionalGeneration(VipLlavaPreTrainedModel, GenerationMixin):
+ def __init__(self, config: VipLlavaConfig):
+ super().__init__(config)
+ self.vision_tower = AutoModel.from_config(config.vision_config)
+
+ self.multi_modal_projector = VipLlavaMultiModalProjector(config)
+ self.vocab_size = config.text_config.vocab_size
+ self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+ if self.language_model._tied_weights_keys is not None:
+ self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys]
+
+ self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+
+ 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 get_output_embeddings(self):
+ return self.language_model.get_output_embeddings()
+
+ def set_output_embeddings(self, new_embeddings):
+ self.language_model.set_output_embeddings(new_embeddings)
+
+ def set_decoder(self, decoder):
+ self.language_model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.language_model.get_decoder()
+
+ # Ignore copy
+ def get_image_features(self, pixel_values: torch.FloatTensor, vision_feature_layers: Union[int, List[int]]):
+ """
+ Obtains image last hidden states from the vision tower and apply multimodal projection.
+
+ Args:
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+ The tensors corresponding to the input images.
+ vision_feature_layers (`Union[int, List[int]]`):
+ The vision feature layer, or the list of indexes of the layers to select
+ the vision feature.
+ Returns:
+ image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+ """
+ image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
+
+ # If multiple feature layers are provided (which is usually the case)
+ # then the image features are concatenated after the CLS is removed.
+ if isinstance(vision_feature_layers, int):
+ image_features = image_outputs.hidden_states[vision_feature_layers][:, 1:]
+ else:
+ # Usually, we select the features from index 1: the layers -2, -5, -8, -11 and 6
+ image_features = [image_outputs.hidden_states[index][:, 1:] for index in vision_feature_layers]
+ image_features = torch.cat(image_features, dim=-1)
+ image_features = self.multi_modal_projector(image_features)
+ return image_features
+
+ @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+ @add_start_docstrings_to_model_forward(VIPLLAVA_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=VipLlavaCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+ # Ignore copy
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = 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,
+ vision_feature_layers: Optional[Union[int, List[int]]] = 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,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **lm_kwargs,
+ ) -> Union[Tuple, VipLlavaCausalLMOutputWithPast]:
+ 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]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import AutoProcessor, VipLlavaForConditionalGeneration
+
+ >>> model = VipLlavaForConditionalGeneration.from_pretrained("llava-hf/vip-llava-7b-hf", device_map="auto", torch_dtype=torch.float16)
+ >>> processor = AutoProcessor.from_pretrained("llava-hf/vip-llava-7b-hf")
+
+ >>> prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.###Human: \n{}###Assistant:"
+ >>> question = "Can you please describe this image?"
+ >>> prompt = prompt.format(question)
+ >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/compel-neg.png"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = processor(text=text, images=image, return_tensors="pt").to(0, torch.float16)
+
+ >>> # Generate
+ >>> generate_ids = model.generate(**inputs, max_new_tokens=20)
+ >>> processor.decode(generate_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+ The image features a brown and white cat sitting on a green surface, with a red ball in its
+ ```"""
+
+ 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
+ vision_feature_layers = (
+ vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers
+ )
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if pixel_values is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values, vision_feature_layers=vision_feature_layers
+ )
+
+ special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
+ special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+ if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
+ n_image_tokens = (input_ids == self.config.image_token_index).sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ 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,
+ logits_to_keep=logits_to_keep,
+ **lm_kwargs,
+ )
+
+ logits = outputs[0]
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ if attention_mask is not None:
+ shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+ shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+ shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+ else:
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = nn.CrossEntropyLoss()
+ loss = loss_fct(
+ shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+ )
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return VipLlavaCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ inputs_embeds=None,
+ pixel_values=None,
+ attention_mask=None,
+ cache_position=None,
+ logits_to_keep=None,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = self.language_model.prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ **kwargs,
+ )
+
+ if cache_position[0] == 0:
+ # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+ # Otherwise we need pixel values to be passed to model
+ model_inputs["pixel_values"] = pixel_values
+
+ return model_inputs
+
+
+__all__ = ["VipLlavaForConditionalGeneration", "VipLlavaPreTrainedModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__init__.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4b68df9c336f8ec77dd6a9efc87175dc6e50eaa6
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_vision_text_dual_encoder import *
+ from .modeling_flax_vision_text_dual_encoder import *
+ from .modeling_tf_vision_text_dual_encoder import *
+ from .modeling_vision_text_dual_encoder import *
+ from .processing_vision_text_dual_encoder import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..498e880055a1717d30b07226d016d11c6b6f9b9a
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/configuration_vision_text_dual_encoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/configuration_vision_text_dual_encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c757bea7a0294d320023fb5312794df00ca1aa7d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/configuration_vision_text_dual_encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_flax_vision_text_dual_encoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_flax_vision_text_dual_encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2ea18d9ec05b9a458036a738711701f223f231d1
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_flax_vision_text_dual_encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_tf_vision_text_dual_encoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_tf_vision_text_dual_encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6fd6616803c4937bd88a4811cffd6492b34ab2a1
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_tf_vision_text_dual_encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_vision_text_dual_encoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_vision_text_dual_encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8c24dc19c3accfd336b675d69c0cafd85402a552
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/modeling_vision_text_dual_encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/processing_vision_text_dual_encoder.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/processing_vision_text_dual_encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..569cf366e7e77dafa66bda2684330c60b641e448
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/__pycache__/processing_vision_text_dual_encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/configuration_vision_text_dual_encoder.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/configuration_vision_text_dual_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..908e1bf1843ed3ff44f13432e203835524c1d9a8
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/configuration_vision_text_dual_encoder.py
@@ -0,0 +1,122 @@
+# 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.
+"""VisionTextDualEncoder model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+from ..auto.configuration_auto import AutoConfig
+from ..chinese_clip.configuration_chinese_clip import ChineseCLIPVisionConfig
+from ..clip.configuration_clip import CLIPVisionConfig
+from ..siglip.configuration_siglip import SiglipVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+VISION_MODEL_CONFIGS = {
+ "clip_vision_model": CLIPVisionConfig,
+ "chinese_clip_vision_model": ChineseCLIPVisionConfig,
+ "siglip_vision_model": SiglipVisionConfig,
+}
+
+
+class VisionTextDualEncoderConfig(PretrainedConfig):
+ r"""
+ [`VisionTextDualEncoderConfig`] is the configuration class to store the configuration of a
+ [`VisionTextDualEncoderModel`]. It is used to instantiate [`VisionTextDualEncoderModel`] model according to the
+ specified arguments, defining the text model and vision model configs.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+ Args:
+ projection_dim (`int`, *optional*, defaults to 512):
+ Dimensionality of text and vision projection layers.
+ logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
+ The initial value of the *logit_scale* parameter. Default is used as per the original CLIP implementation.
+ kwargs (*optional*):
+ Dictionary of keyword arguments.
+
+ Examples:
+
+ ```python
+ >>> from transformers import ViTConfig, BertConfig, VisionTextDualEncoderConfig, VisionTextDualEncoderModel
+
+ >>> # Initializing a BERT and ViT configuration
+ >>> config_vision = ViTConfig()
+ >>> config_text = BertConfig()
+
+ >>> config = VisionTextDualEncoderConfig.from_vision_text_configs(config_vision, config_text, projection_dim=512)
+
+ >>> # Initializing a BERT and ViT model (with random weights)
+ >>> model = VisionTextDualEncoderModel(config=config)
+
+ >>> # Accessing the model configuration
+ >>> config_vision = model.config.vision_config
+ >>> config_text = model.config.text_config
+
+ >>> # Saving the model, including its configuration
+ >>> model.save_pretrained("vit-bert")
+
+ >>> # loading model and config from pretrained folder
+ >>> vision_text_config = VisionTextDualEncoderConfig.from_pretrained("vit-bert")
+ >>> model = VisionTextDualEncoderModel.from_pretrained("vit-bert", config=vision_text_config)
+ ```"""
+
+ model_type = "vision-text-dual-encoder"
+ sub_configs = {"vision_config": AutoConfig, "text_config": AutoConfig}
+ is_composition = True
+
+ def __init__(self, projection_dim=512, logit_scale_init_value=2.6592, **kwargs):
+ super().__init__(**kwargs)
+
+ if "vision_config" not in kwargs:
+ raise ValueError("`vision_config` can not be `None`.")
+
+ if "text_config" not in kwargs:
+ raise ValueError("`text_config` can not be `None`.")
+
+ vision_config = kwargs.pop("vision_config")
+ text_config = kwargs.pop("text_config")
+
+ vision_model_type = vision_config.pop("model_type")
+ text_model_type = text_config.pop("model_type")
+
+ vision_config_class = VISION_MODEL_CONFIGS.get(vision_model_type)
+ if vision_config_class is not None:
+ self.vision_config = vision_config_class(**vision_config)
+ else:
+ self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config)
+ if hasattr(self.vision_config, "vision_config"):
+ self.vision_config = self.vision_config.vision_config
+
+ self.text_config = AutoConfig.for_model(text_model_type, **text_config)
+
+ self.projection_dim = projection_dim
+ self.logit_scale_init_value = logit_scale_init_value
+
+ @classmethod
+ def from_vision_text_configs(cls, vision_config: PretrainedConfig, text_config: PretrainedConfig, **kwargs):
+ r"""
+ Instantiate a [`VisionTextDualEncoderConfig`] (or a derived class) from text model configuration and vision
+ model configuration.
+
+ Returns:
+ [`VisionTextDualEncoderConfig`]: An instance of a configuration object
+ """
+
+ return cls(vision_config=vision_config.to_dict(), text_config=text_config.to_dict(), **kwargs)
+
+
+__all__ = ["VisionTextDualEncoderConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_flax_vision_text_dual_encoder.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_flax_vision_text_dual_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..b12327d8ca0253a8ce1a5217e23c2d6c05600906
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_flax_vision_text_dual_encoder.py
@@ -0,0 +1,601 @@
+# 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.
+"""Flax VisionTextDualEncoder model."""
+
+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.traverse_util import flatten_dict, unflatten_dict
+
+from ...modeling_flax_utils import FlaxPreTrainedModel, append_replace_return_docstrings, overwrite_call_docstring
+from ...utils import add_start_docstrings, logging
+from ..auto.configuration_auto import AutoConfig
+from ..auto.modeling_flax_auto import FLAX_MODEL_MAPPING, FlaxAutoModel
+from ..clip.modeling_flax_clip import FlaxCLIPOutput, FlaxCLIPVisionModel
+from .configuration_vision_text_dual_encoder import VisionTextDualEncoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "VisionTextDualEncoderConfig"
+
+VISION_TEXT_DUAL_ENCODER_START_DOCSTRING = r"""
+ This class can be used to initialize a vision-text dual encoder model with any pretrained vision autoencoding model
+ as the vision encoder and any pretrained text model as the text encoder. The vision and text encoders are loaded
+ via the [`~FlaxAutoModel.from_pretrained`] method. The projection layers are automatically added to the model and
+ should be fine-tuned on a downstream task, like contrastive image-text modeling.
+
+ In [LiT: Zero-Shot Transfer with Locked-image Text Tuning](https://arxiv.org/abs/2111.07991) it is shown how
+ leveraging pre-trained (locked/frozen) image and text model for contrastive learning yields significant improvment
+ on new zero-shot vision tasks such as image classification or retrieval.
+
+ After such a Vision-Text-Dual-Encoder 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
+ [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 ([`VisionTextDualEncoderConfig`]): 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`].
+"""
+
+
+VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`numpy.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 [`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)
+ 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]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ 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
+ an image processor (e.g. if you use ViT as the encoder, you should use [`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.
+"""
+
+
+class FlaxVisionTextDualEncoderModule(nn.Module):
+ config: VisionTextDualEncoderConfig
+ dtype: jnp.dtype = jnp.float32
+
+ def setup(self):
+ vision_config = self.config.vision_config
+ text_config = self.config.text_config
+
+ self.vision_embed_dim = vision_config.hidden_size
+ self.text_embed_dim = text_config.hidden_size
+ self.projection_dim = self.config.projection_dim
+
+ vision_module = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__, FlaxCLIPVisionModel).module_class
+ text_module = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
+
+ self.vision_model = vision_module(vision_config, dtype=self.dtype)
+ self.text_model = text_module(text_config, dtype=self.dtype)
+
+ self.visual_projection = nn.Dense(
+ self.projection_dim,
+ dtype=self.dtype,
+ kernel_init=jax.nn.initializers.normal(0.02),
+ use_bias=False,
+ )
+ self.text_projection = nn.Dense(
+ self.projection_dim,
+ dtype=self.dtype,
+ kernel_init=jax.nn.initializers.normal(0.02),
+ use_bias=False,
+ )
+
+ self.logit_scale = self.param(
+ "logit_scale", lambda _, shape: jnp.ones(shape) * self.config.logit_scale_init_value, []
+ )
+
+ def __call__(
+ self,
+ input_ids=None,
+ pixel_values=None,
+ attention_mask=None,
+ position_ids=None,
+ token_type_ids=None,
+ deterministic: bool = True,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ deterministic=deterministic,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ text_outputs = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ deterministic=deterministic,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ image_embeds = vision_outputs[1]
+ image_embeds = self.visual_projection(image_embeds)
+
+ text_embeds = text_outputs[1]
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / jnp.linalg.norm(image_embeds, axis=-1, keepdims=True)
+ text_embeds = text_embeds / jnp.linalg.norm(text_embeds, axis=-1, keepdims=True)
+
+ # cosine similarity as logits
+ logit_scale = jnp.exp(self.logit_scale)
+ logits_per_text = jnp.matmul(text_embeds, image_embeds.T) * logit_scale
+ logits_per_image = logits_per_text.T
+
+ if not return_dict:
+ return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+
+ return FlaxCLIPOutput(
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+@add_start_docstrings(VISION_TEXT_DUAL_ENCODER_START_DOCSTRING)
+class FlaxVisionTextDualEncoderModel(FlaxPreTrainedModel):
+ config_class = VisionTextDualEncoderConfig
+ module_class = FlaxVisionTextDualEncoderModule
+
+ def __init__(
+ self,
+ config: VisionTextDualEncoderConfig,
+ input_shape: Optional[Tuple] = None,
+ seed: int = 0,
+ dtype: jnp.dtype = jnp.float32,
+ _do_init: bool = True,
+ **kwargs,
+ ):
+ if not _do_init:
+ raise ValueError(
+ "`FlaxVisionTextDualEncoderModel` cannot be created without initializing, `_do_init` must be `True`."
+ )
+
+ if input_shape is None:
+ input_shape = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))
+
+ module = self.module_class(config=config, dtype=dtype, **kwargs)
+ super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+
+ def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
+ # init input tensor
+ input_ids = jnp.zeros(input_shape[0], dtype="i4")
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape[0])
+ token_type_ids = jnp.ones_like(input_ids)
+ attention_mask = jnp.ones_like(input_ids)
+
+ pixel_values = jax.random.normal(rng, input_shape[1])
+
+ params_rng, dropout_rng = jax.random.split(rng)
+ rngs = {"params": params_rng, "dropout": dropout_rng}
+
+ random_params = self.module.init(rngs, input_ids, pixel_values, attention_mask, position_ids, token_type_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 __call__(
+ self,
+ input_ids,
+ pixel_values,
+ attention_mask=None,
+ position_ids=None,
+ token_type_ids=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
+
+ pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))
+
+ if position_ids is None:
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
+
+ if token_type_ids is None:
+ token_type_ids = jnp.zeros_like(input_ids)
+
+ 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(pixel_values, dtype=jnp.float32),
+ jnp.array(attention_mask, dtype="i4"),
+ jnp.array(position_ids, dtype="i4"),
+ jnp.array(token_type_ids, dtype="i4"),
+ not train,
+ output_attentions,
+ output_hidden_states,
+ return_dict,
+ rngs=rngs,
+ )
+
+ def get_text_features(
+ self,
+ input_ids,
+ attention_mask=None,
+ position_ids=None,
+ token_type_ids=None,
+ params: dict = None,
+ dropout_rng: jax.random.PRNGKey = None,
+ train=False,
+ ):
+ r"""
+ Args:
+ input_ids (`numpy.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)
+
+ Returns:
+ text_features (`jnp.ndarray` of shape `(batch_size, output_dim`): The text embeddings obtained by applying
+ the projection layer to the pooled output of text model.
+ """
+ if position_ids is None:
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
+
+ if token_type_ids is None:
+ token_type_ids = jnp.zeros_like(input_ids)
+
+ 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
+
+ def _get_features(module, input_ids, attention_mask, position_ids, token_type_ids, deterministic):
+ text_outputs = module.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ deterministic=deterministic,
+ )
+ pooled_output = text_outputs[1]
+ text_features = module.text_projection(pooled_output)
+ return text_features
+
+ return self.module.apply(
+ {"params": params or self.params},
+ jnp.array(input_ids, dtype="i4"),
+ jnp.array(attention_mask, dtype="i4"),
+ jnp.array(position_ids, dtype="i4"),
+ jnp.array(token_type_ids, dtype="i4"),
+ not train,
+ method=_get_features,
+ rngs=rngs,
+ )
+
+ def get_image_features(
+ self, pixel_values, params: dict = None, dropout_rng: jax.random.PRNGKey = None, train=False
+ ):
+ r"""
+ Args:
+ pixel_values (`numpy.ndarray` 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 [`ImageFeatureExtractionMixin`]. See [`ImageFeatureExtractionMixin.__call__`] for details.
+
+ Returns:
+ image_features (`jnp.ndarray` of shape `(batch_size, output_dim`): The image embeddings obtained by
+ applying the projection layer to the pooled output of vision model.
+ """
+
+ # Handle any PRNG if needed
+ rngs = {}
+ if dropout_rng is not None:
+ rngs["dropout"] = dropout_rng
+
+ def _get_features(module, pixel_values, deterministic):
+ vision_outputs = module.vision_model(pixel_values=pixel_values, deterministic=deterministic)
+ pooled_output = vision_outputs[1] # pooled_output
+ image_features = module.visual_projection(pooled_output)
+ return image_features
+
+ return self.module.apply(
+ {"params": params or self.params},
+ jnp.array(pixel_values, dtype=jnp.float32),
+ not train,
+ method=_get_features,
+ rngs=rngs,
+ )
+
+ @classmethod
+ def from_vision_text_pretrained(
+ cls,
+ vision_model_name_or_path: Optional[str] = None,
+ text_model_name_or_path: Optional[str] = None,
+ *model_args,
+ **kwargs,
+ ) -> FlaxPreTrainedModel:
+ """
+ Params:
+ vision_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the vision model. 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/`.
+ - A path or url to a *PyTorch checkpoint folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument. This
+ loading path is slower than converting the PyTorch checkpoint in a Flax model using the provided
+ conversion scripts and loading the Flax model afterwards.
+
+ text_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text model. 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/`.
+ - A path or url to a *PyTorch checkpoint folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument. This
+ loading path is slower than converting the PyTorch checkpoint in a Flax model using the provided
+ conversion scripts and loading the Flax model afterwards.
+
+ 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 text configuration, use the prefix *text_* for each configuration parameter.
+ - To update the vision configuration, use the prefix *vision_* 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 FlaxVisionTextDualEncoderModel
+
+ >>> # initialize a model from pretrained ViT and BERT models. Note that the projection layers will be randomly initialized.
+ >>> model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./vit-bert")
+ >>> # load fine-tuned model
+ >>> model = FlaxVisionTextDualEncoderModel.from_pretrained("./vit-bert")
+ ```"""
+
+ kwargs_vision = {
+ argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
+ }
+
+ kwargs_text = {
+ argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
+ }
+
+ # remove text, vision kwargs from kwargs
+ for key in kwargs_vision.keys():
+ del kwargs["vision_" + key]
+ for key in kwargs_text.keys():
+ del kwargs["text_" + key]
+
+ # Load and initialize the text and vision model
+ vision_model = kwargs_vision.pop("model", None)
+ if vision_model is None:
+ if vision_model_name_or_path is None:
+ raise ValueError(
+ "If `vision_model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
+ )
+
+ if "config" not in kwargs_vision:
+ vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)
+
+ if vision_config.model_type == "clip":
+ kwargs_vision["config"] = vision_config.vision_config
+ vision_model = FlaxCLIPVisionModel.from_pretrained(
+ vision_model_name_or_path, *model_args, **kwargs_vision
+ )
+ else:
+ kwargs_vision["config"] = vision_config
+ vision_model = FlaxAutoModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)
+
+ text_model = kwargs_text.pop("model", None)
+ if text_model is None:
+ if text_model_name_or_path is None:
+ raise ValueError(
+ "If `text_model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
+ )
+
+ if "config" not in kwargs_text:
+ text_config = AutoConfig.from_pretrained(text_model_name_or_path)
+ kwargs_text["config"] = text_config
+
+ text_model = FlaxAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)
+
+ # instantiate config with corresponding kwargs
+ dtype = kwargs.pop("dtype", jnp.float32)
+ config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config, **kwargs)
+
+ # init model
+ model = cls(config, *model_args, dtype=dtype, **kwargs)
+
+ model.params["vision_model"] = vision_model.params
+ model.params["text_model"] = text_model.params
+
+ # the projection layers are always newly initialized when loading the model
+ # using pre-trained vision and text model.
+ logger.warning(
+ "The projection layer and logit scale weights `[('visual_projection', 'kernel'), ('text_projection',"
+ " 'kernel'), ('logit_scale',)]` are newly initialized. You should probably TRAIN this model on a"
+ " down-stream task to be able to use it for predictions and inference."
+ )
+
+ return model
+
+
+VISION_TEXT_DUAL_ENCODER_MODEL_DOCSTRING = r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> import jax
+ >>> from transformers import (
+ ... FlaxVisionTextDualEncoderModel,
+ ... VisionTextDualEncoderProcessor,
+ ... AutoImageProcessor,
+ ... AutoTokenizer,
+ ... )
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ >>> image_processor = AutoImageProcesor.from_pretrained("google/vit-base-patch16-224")
+ >>> processor = VisionTextDualEncoderProcessor(image_processor, tokenizer)
+ >>> model = FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+
+ >>> # contrastive training
+ >>> urls = [
+ ... "http://images.cocodataset.org/val2017/000000039769.jpg",
+ ... "https://farm3.staticflickr.com/2674/5850229113_4fe05d5265_z.jpg",
+ ... ]
+ >>> images = [Image.open(requests.get(url, stream=True).raw) for url in urls]
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=images, return_tensors="np", padding=True
+ ... )
+ >>> outputs = model(
+ ... input_ids=inputs.input_ids,
+ ... attention_mask=inputs.attention_mask,
+ ... pixel_values=inputs.pixel_values,
+ ... )
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+
+ >>> # save and load from pretrained
+ >>> model.save_pretrained("vit-bert")
+ >>> model = FlaxVisionTextDualEncoderModel.from_pretrained("vit-bert")
+
+ >>> # inference
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = jax.nn.softmax(logits_per_image, axis=1) # we can take the softmax to get the label probabilities
+ ```
+"""
+
+overwrite_call_docstring(
+ FlaxVisionTextDualEncoderModel,
+ VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING + VISION_TEXT_DUAL_ENCODER_MODEL_DOCSTRING,
+)
+append_replace_return_docstrings(
+ FlaxVisionTextDualEncoderModel, output_type=FlaxCLIPOutput, config_class=_CONFIG_FOR_DOC
+)
+
+
+__all__ = ["FlaxVisionTextDualEncoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_tf_vision_text_dual_encoder.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_tf_vision_text_dual_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca88d2fec95ef8f5b3933d888f6cf76ff5c4770d
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_tf_vision_text_dual_encoder.py
@@ -0,0 +1,624 @@
+# 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.
+"""TensorFlow VisionTextDualEncoder model."""
+
+from __future__ import annotations
+
+import re
+from typing import Optional, Tuple, Union
+
+import tensorflow as tf
+
+from ...configuration_utils import PretrainedConfig
+from ...modeling_tf_utils import TFPreTrainedModel, keras, unpack_inputs
+from ...tf_utils import shape_list
+from ...utils import (
+ DUMMY_INPUTS,
+ 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
+from ..clip.modeling_tf_clip import CLIPVisionConfig, TFCLIPOutput, TFCLIPVisionModel
+from .configuration_vision_text_dual_encoder import VisionTextDualEncoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "VisionTextDualEncoderConfig"
+
+VISION_TEXT_DUAL_ENCODER_START_DOCSTRING = r"""
+ This class can be used to initialize a vision-text dual encoder model with any pretrained vision autoencoding model
+ as the vision encoder and any pretrained text model as the text encoder. The vision and text encoders are loaded
+ via the [`~TFAutoModel.from_pretrained`] method. The projection layers are automatically added to the model and
+ should be fine-tuned on a downstream task, like contrastive image-text modeling.
+
+ In [LiT: Zero-Shot Transfer with Locked-image Text Tuning](https://arxiv.org/abs/2111.07991) it is shown how
+ leveraging pre-trained (locked/frozen) image and text model for contrastive learning yields significant improvment
+ on new zero-shot vision tasks such as image classification or retrieval.
+
+ After such a Vision-Text-Dual-Encoder 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 Keras Model and refer to the TF documentation for all matter related to general usage and behavior.
+
+ Parameters:
+ config ([`VisionEncoderDecoderConfig`]): 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.
+"""
+
+
+VISION_TEXT_DUAL_ENCODER_TEXT_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`tf.Tensor` 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 (`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)
+ position_ids (`tf.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)
+ 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.
+"""
+
+VISION_TEXT_DUAL_ENCODER_VISION_INPUTS_DOCSTRING = r"""
+ Args:
+ pixel_values (`tf.Tensor` 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.
+"""
+
+VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`tf.Tensor` 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 (`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)
+ position_ids (`tf.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)
+ pixel_values (`tf.Tensor` 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
+ an image processor (e.g. if you use ViT as the encoder, you should use [`AutoImageProcessor`]). See
+ [`ViTImageProcessor.__call__`] for details.
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+ 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_tf_clip.contrastive_loss
+def contrastive_loss(logits: tf.Tensor) -> tf.Tensor:
+ return tf.math.reduce_mean(
+ keras.metrics.sparse_categorical_crossentropy(
+ y_true=tf.range(shape_list(logits)[0]), y_pred=logits, from_logits=True
+ )
+ )
+
+
+# Copied from transformers.models.clip.modeling_tf_clip.clip_loss
+def clip_loss(similarity: tf.Tensor) -> tf.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(tf.transpose(similarity))
+ return (caption_loss + image_loss) / 2.0
+
+
+@add_start_docstrings(VISION_TEXT_DUAL_ENCODER_START_DOCSTRING)
+class TFVisionTextDualEncoderModel(TFPreTrainedModel):
+ config_class = VisionTextDualEncoderConfig
+ base_model_prefix = "vision_text_dual_encoder"
+ load_weight_prefix = "tf_vision_text_dual_encoder_model"
+
+ def __init__(
+ self,
+ config: Optional[VisionTextDualEncoderConfig] = None,
+ vision_model: Optional[TFPreTrainedModel] = None,
+ text_model: Optional[TFPreTrainedModel] = None,
+ ):
+ if config is None and (vision_model is None or text_model is None):
+ raise ValueError("Either a configuration or an vision and a text model has to be provided")
+
+ if config is None:
+ config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config)
+ else:
+ if not isinstance(config, self.config_class):
+ raise ValueError(f"config: {config} has to be of type {self.config_class}")
+
+ # initialize with config
+ super().__init__(config)
+
+ if vision_model is None:
+ if isinstance(config.vision_config, CLIPVisionConfig):
+ vision_model = TFCLIPVisionModel.from_config(config.vision_config, name="vision_model")
+ else:
+ vision_model = TFAutoModel.from_config(config.vision_config, name="vision_model")
+
+ if text_model is None:
+ text_model = TFAutoModel.from_config(config.text_config, name="text_model")
+
+ self.vision_model = vision_model
+ self.text_model = text_model
+
+ # make sure that the individual model's config refers to the shared config
+ # so that the updates to the config will be synced
+ self.vision_model.config = self.config.vision_config
+ self.text_model.config = self.config.text_config
+
+ self.vision_embed_dim = config.vision_config.hidden_size
+ self.text_embed_dim = config.text_config.hidden_size
+ self.projection_dim = config.projection_dim
+
+ self.visual_projection = keras.layers.Dense(self.projection_dim, use_bias=False, name="visual_projection")
+ self.text_projection = keras.layers.Dense(self.projection_dim, use_bias=False, name="text_projection")
+ self.logit_scale = None
+ self.config = config
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ # Build in the build() method to make sure the names are right
+ initializer = keras.initializers.Constant(self.config.logit_scale_init_value)
+ self.logit_scale = self.add_weight(shape=(1,), initializer=initializer, name="logit_scale")
+
+ if getattr(self, "visual_projection", None) is not None:
+ with tf.name_scope(self.visual_projection.name):
+ self.visual_projection.build([None, None, self.vision_embed_dim])
+ if getattr(self, "text_projection", None) is not None:
+ with tf.name_scope(self.text_projection.name):
+ self.text_projection.build([None, None, self.text_embed_dim])
+ with tf.name_scope(self.vision_model.name):
+ self.vision_model.build(None)
+ with tf.name_scope(self.text_model.name):
+ self.text_model.build(None)
+
+ 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.
+ if "vision_model" in tf_weight:
+ if tf_weight.count("vision_model") == 1:
+ return (re.sub(r"vision_model\..*?\.", "vision_model.", tf_weight),)
+ elif tf_weight.count("vision_model") == 2:
+ return (re.sub(r"vision_model\..*?\.vision_model", "vision_model.vision_model", tf_weight),)
+ else:
+ raise ValueError(
+ f"Unexpected weight name {tf_weight}. Please file an issue on the"
+ " Transformers repo to let us know about this error!"
+ )
+ elif "text_model" in tf_weight:
+ return (re.sub(r"text_model\..*?\.", "text_model.", tf_weight),)
+ else:
+ return (tf_weight,)
+
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_TEXT_INPUTS_DOCSTRING)
+ def get_text_features(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ token_type_ids=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ r"""
+ Returns:
+ text_features (`tf.Tensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying
+ the projection layer to the pooled output of [`TFCLIPTextModel`].
+
+ Examples:
+
+ ```python
+ >>> from transformers import TFVisionTextDualEncoderModel, AutoTokenizer
+
+ >>> model = TFVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian", from_pt=True)
+ >>> tokenizer = AutoTokenizer.from_pretrained("clip-italian/clip-italian")
+
+ >>> inputs = tokenizer(["una foto di un gatto", "una foto di un cane"], padding=True, return_tensors="np")
+ >>> text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs = 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,
+ )
+
+ pooled_output = text_outputs[1]
+ text_features = self.text_projection(pooled_output)
+
+ return text_features
+
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_VISION_INPUTS_DOCSTRING)
+ def get_image_features(
+ self,
+ pixel_values=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ r"""
+ Returns:
+ image_features (`tf.Tensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying
+ the projection layer to the pooled output of [`TFCLIPVisionModel`].
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import TFVisionTextDualEncoderModel, AutoImageProcessor
+
+ >>> model = TFVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian", from_pt=True)
+ >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = image_processor(images=image, return_tensors="np")
+
+ >>> image_features = model.get_image_features(**inputs)
+ ```"""
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = vision_outputs[1] # pooled_output
+ image_features = self.visual_projection(pooled_output)
+
+ return image_features
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=TFCLIPOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: tf.Tensor | None = None,
+ pixel_values: tf.Tensor | None = None,
+ attention_mask: tf.Tensor | None = None,
+ position_ids: tf.Tensor | None = None,
+ return_loss: Optional[bool] = None,
+ token_type_ids: 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], TFCLIPOutput]:
+ r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import (
+ ... TFVisionTextDualEncoderModel,
+ ... VisionTextDualEncoderProcessor,
+ ... AutoImageProcessor,
+ ... AutoTokenizer,
+ ... )
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+ >>> processor = VisionTextDualEncoderProcessor(image_processor, tokenizer)
+ >>> model = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+
+ >>> # contrastive training
+ >>> urls = [
+ ... "http://images.cocodataset.org/val2017/000000039769.jpg",
+ ... "https://farm3.staticflickr.com/2674/5850229113_4fe05d5265_z.jpg",
+ ... ]
+ >>> images = [Image.open(requests.get(url, stream=True).raw) for url in urls]
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=images, return_tensors="np", padding=True
+ ... )
+ >>> outputs = model(
+ ... input_ids=inputs.input_ids,
+ ... attention_mask=inputs.attention_mask,
+ ... pixel_values=inputs.pixel_values,
+ ... return_loss=True,
+ ... )
+ >>> loss, logits_per_image = outputs.loss, outputs.logits_per_image # this is the image-text similarity score
+
+ >>> # save and load from pretrained
+ >>> model.save_pretrained("vit-bert")
+ >>> model = TFVisionTextDualEncoderModel.from_pretrained("vit-bert")
+
+ >>> # inference
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = tf.nn.softmax(logits_per_image, axis=1) # we can take the softmax to get the label probabilities
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ training=training,
+ )
+
+ 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,
+ training=training,
+ )
+
+ image_embeds = vision_outputs[1] # pooler_output
+ image_embeds = self.visual_projection(image_embeds)
+
+ text_embeds = text_outputs[1] # pooler_output
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / tf.norm(image_embeds, axis=-1, keepdims=True)
+ text_embeds = text_embeds / tf.norm(text_embeds, axis=-1, keepdims=True)
+
+ # cosine similarity as logits
+ logit_scale = tf.math.exp(self.logit_scale)
+ logits_per_text = tf.matmul(text_embeds, image_embeds, transpose_b=True) * logit_scale
+ logits_per_image = tf.transpose(logits_per_text)
+
+ loss = None
+ if return_loss:
+ loss = clip_loss(logits_per_text)
+ if loss.shape.rank == 0:
+ loss = tf.expand_dims(loss, 0)
+
+ if not return_dict:
+ output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+ return ((loss,) + output) if loss is not None else output
+
+ return TFCLIPOutput(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+ @classmethod
+ def from_vision_text_pretrained(
+ cls,
+ vision_model_name_or_path: Optional[str] = None,
+ text_model_name_or_path: Optional[str] = None,
+ *model_args,
+ **kwargs,
+ ) -> TFPreTrainedModel:
+ """
+ Params:
+ vision_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the vision model. 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 folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument.
+
+ text_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text model. 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 folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument.
+
+ 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 text configuration, use the prefix *text_* for each configuration parameter.
+ - To update the vision configuration, use the prefix *vision_* 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 TFVisionTextDualEncoderModel
+
+ >>> # initialize a model from pretrained ViT and BERT models. Note that the projection layers will be randomly initialized.
+ >>> model = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./vit-bert")
+ >>> # load fine-tuned model
+ >>> model = TFVisionTextDualEncoderModel.from_pretrained("./vit-bert")
+ ```"""
+ kwargs_vision = {
+ argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
+ }
+
+ kwargs_text = {
+ argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
+ }
+
+ # remove vision, text kwargs from kwargs
+ for key in kwargs_vision.keys():
+ del kwargs["vision_" + key]
+ for key in kwargs_text.keys():
+ del kwargs["text_" + key]
+
+ # Load and initialize the vision and text model
+ vision_model = kwargs_vision.pop("model", None)
+ if vision_model is None:
+ if vision_model_name_or_path is None:
+ raise ValueError(
+ "If `vision_model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
+ )
+ kwargs_vision["name"] = "vision_model"
+ kwargs_vision["load_weight_prefix"] = cls.load_weight_prefix
+
+ vision_config_dict, unused_args = PretrainedConfig.get_config_dict(vision_model_name_or_path, **kwargs)
+ if vision_config_dict.get("model_type", None) == "clip_vision_model":
+ vision_config = CLIPVisionConfig.from_dict(vision_config_dict)
+ else:
+ vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)
+
+ if vision_config.model_type == "clip_vision_model":
+ kwargs_vision["config"] = vision_config
+ vision_class = TFCLIPVisionModel
+ elif vision_config.model_type == "clip":
+ kwargs_vision["config"] = vision_config.vision_config
+ vision_class = TFCLIPVisionModel
+ else:
+ kwargs_vision["config"] = vision_config
+ vision_class = TFAutoModel
+ vision_model = vision_class.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)
+
+ text_model = kwargs_text.pop("model", None)
+ if text_model is None:
+ if text_model_name_or_path is None:
+ raise ValueError(
+ "If `text_model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
+ )
+ kwargs_text["name"] = "text_model"
+ kwargs_text["load_weight_prefix"] = cls.load_weight_prefix
+
+ if "config" not in kwargs_text:
+ text_config = AutoConfig.from_pretrained(text_model_name_or_path)
+ kwargs_text["config"] = text_config
+
+ text_model = TFAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)
+
+ # instantiate config with corresponding kwargs
+ config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config, **kwargs)
+
+ # init model
+ model = cls(config=config, vision_model=vision_model, text_model=text_model)
+
+ # the projection layers are always newly initialized when loading the model
+ # using pre-trained vision and text model.
+ logger.warning(
+ "The projection layer and logit scale weights `['visual_projection.weight', 'text_projection.weight',"
+ " 'logit_scale']` are newly initialized. You should probably TRAIN this model on a down-stream task to be"
+ " able to use it for predictions and inference."
+ )
+
+ if vision_model.name != "vision_model":
+ raise ValueError("vision model must be created with the name `vision_model`.")
+ if text_model.name != "text_model":
+ raise ValueError("text model must be created with the name `text_model`.")
+
+ model.build_in_name_scope() # Ensure model is fully built
+
+ return model
+
+ @property
+ def dummy_inputs(self):
+ """
+ Dummy inputs to build the network.
+
+ Returns:
+ `Dict[str, tf.Tensor]`: The dummy inputs.
+ """
+ input_ids = tf.constant(DUMMY_INPUTS, dtype=tf.int32)
+ batch_size, seq_len = input_ids.shape
+
+ VISION_DUMMY_INPUTS = tf.random.uniform(
+ shape=(
+ batch_size,
+ self.config.vision_config.num_channels,
+ self.config.vision_config.image_size,
+ self.config.vision_config.image_size,
+ ),
+ dtype=tf.float32,
+ )
+ pixel_values = tf.constant(VISION_DUMMY_INPUTS)
+ dummy = {"pixel_values": pixel_values, "input_ids": input_ids}
+ return dummy
+
+
+__all__ = ["TFVisionTextDualEncoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e770f6935d86e37a1a6b0579940142a411da0da
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py
@@ -0,0 +1,532 @@
+# 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 VisionTextDualEncoder model."""
+
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+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
+from ..clip.modeling_clip import CLIPOutput, CLIPVisionConfig, CLIPVisionModel
+from .configuration_vision_text_dual_encoder import VisionTextDualEncoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "VisionTextDualEncoderConfig"
+
+VISION_TEXT_DUAL_ENCODER_START_DOCSTRING = r"""
+ This class can be used to initialize a vision-text dual encoder model with any pretrained vision autoencoding model
+ as the vision encoder and any pretrained text model as the text encoder. The vision and text encoders are loaded
+ via the [`~AutoModel.from_pretrained`] method. The projection layers are automatically added to the model and
+ should be fine-tuned on a downstream task, like contrastive image-text modeling.
+
+ In [LiT: Zero-Shot Transfer with Locked-image Text Tuning](https://arxiv.org/abs/2111.07991) it is shown how
+ leveraging pre-trained (locked/frozen) image and text model for contrastive learning yields significant improvment
+ on new zero-shot vision tasks such as image classification or retrieval.
+
+ After such a Vision-Text-Dual-Encoder 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 ([`VisionEncoderDecoderConfig`]): 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.
+"""
+
+
+VISION_TEXT_DUAL_ENCODER_TEXT_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 [`PreTrainedTokenizer`]. 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)
+ 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)
+ 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.
+"""
+
+VISION_TEXT_DUAL_ENCODER_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.
+"""
+
+VISION_TEXT_DUAL_ENCODER_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)
+ 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)
+ 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
+ an image processor (e.g. if you use ViT as the encoder, you should use [`AutoImageProcessor`]). See
+ [`ViTImageProcessor.__call__`] for details.
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+ 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.contrastive_loss
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+# Copied from transformers.models.clip.modeling_clip.clip_loss
+def clip_loss(similarity: torch.Tensor) -> torch.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(similarity.t())
+ return (caption_loss + image_loss) / 2.0
+
+
+@add_start_docstrings(VISION_TEXT_DUAL_ENCODER_START_DOCSTRING)
+class VisionTextDualEncoderModel(PreTrainedModel):
+ config_class = VisionTextDualEncoderConfig
+ base_model_prefix = "vision_text_dual_encoder"
+ _supports_flash_attn_2 = True
+ _supports_sdpa = True
+
+ def __init__(
+ self,
+ config: Optional[VisionTextDualEncoderConfig] = None,
+ vision_model: Optional[PreTrainedModel] = None,
+ text_model: Optional[PreTrainedModel] = None,
+ ):
+ if config is None and (vision_model is None or text_model is None):
+ raise ValueError("Either a configuration or an vision and a text model has to be provided")
+
+ if config is None:
+ config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config)
+ else:
+ if not isinstance(config, self.config_class):
+ raise ValueError(f"config: {config} has to be of type {self.config_class}")
+
+ # initialize with config
+ super().__init__(config)
+
+ if vision_model is None:
+ if isinstance(config.vision_config, CLIPVisionConfig):
+ vision_model = CLIPVisionModel(config.vision_config)
+ else:
+ vision_model = AutoModel.from_config(config.vision_config)
+
+ if text_model is None:
+ text_model = AutoModel.from_config(config.text_config)
+
+ self.vision_model = vision_model
+ self.text_model = text_model
+
+ # make sure that the individual model's config refers to the shared config
+ # so that the updates to the config will be synced
+ self.config.vision_config._attn_implementation = self.vision_model.config._attn_implementation
+ self.config.text_config._attn_implementation = self.text_model.config._attn_implementation
+ self.vision_model.config = self.config.vision_config
+ self.text_model.config = self.config.text_config
+
+ self.vision_embed_dim = config.vision_config.hidden_size
+ self.text_embed_dim = config.text_config.hidden_size
+ self.projection_dim = config.projection_dim
+
+ self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
+ self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_TEXT_INPUTS_DOCSTRING)
+ def get_text_features(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ token_type_ids=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ 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 [`CLIPTextModel`].
+
+ Examples:
+
+ ```python
+ >>> from transformers import VisionTextDualEncoderModel, AutoTokenizer
+
+ >>> model = VisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian")
+ >>> tokenizer = AutoTokenizer.from_pretrained("clip-italian/clip-italian")
+
+ >>> inputs = tokenizer(["una foto di un gatto", "una foto di un cane"], padding=True, return_tensors="pt")
+ >>> text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs = 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,
+ )
+
+ pooled_output = text_outputs[1]
+ text_features = self.text_projection(pooled_output)
+
+ return text_features
+
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_VISION_INPUTS_DOCSTRING)
+ def get_image_features(
+ self,
+ pixel_values=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ 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 [`CLIPVisionModel`].
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import VisionTextDualEncoderModel, AutoImageProcessor
+
+ >>> model = VisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian")
+ >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = Image.open(requests.get(url, stream=True).raw)
+
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> image_features = model.get_image_features(**inputs)
+ ```"""
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = vision_outputs[1] # pooled_output
+ image_features = self.visual_projection(pooled_output)
+
+ return image_features
+
+ @add_start_docstrings_to_model_forward(VISION_TEXT_DUAL_ENCODER_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=CLIPOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ return_loss: Optional[bool] = None,
+ token_type_ids: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple[torch.Tensor], CLIPOutput]:
+ r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import requests
+ >>> from transformers import (
+ ... VisionTextDualEncoderModel,
+ ... VisionTextDualEncoderProcessor,
+ ... AutoImageProcessor,
+ ... AutoTokenizer,
+ ... )
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+ >>> processor = VisionTextDualEncoderProcessor(image_processor, tokenizer)
+ >>> model = VisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+
+ >>> # contrastive training
+ >>> urls = [
+ ... "http://images.cocodataset.org/val2017/000000039769.jpg",
+ ... "https://farm3.staticflickr.com/2674/5850229113_4fe05d5265_z.jpg",
+ ... ]
+ >>> images = [Image.open(requests.get(url, stream=True).raw) for url in urls]
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=images, return_tensors="pt", padding=True
+ ... )
+ >>> outputs = model(
+ ... input_ids=inputs.input_ids,
+ ... attention_mask=inputs.attention_mask,
+ ... pixel_values=inputs.pixel_values,
+ ... return_loss=True,
+ ... )
+ >>> loss, logits_per_image = outputs.loss, outputs.logits_per_image # this is the image-text similarity score
+
+ >>> # save and load from pretrained
+ >>> model.save_pretrained("vit-bert")
+ >>> model = VisionTextDualEncoderModel.from_pretrained("vit-bert")
+
+ >>> # inference
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ 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,
+ )
+
+ image_embeds = vision_outputs[1] # pooler_output
+ image_embeds = self.visual_projection(image_embeds)
+
+ text_embeds = text_outputs[1] # pooler_output
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
+ text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale = self.logit_scale.exp()
+ logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
+ logits_per_image = logits_per_text.T
+
+ loss = None
+ if return_loss:
+ loss = clip_loss(logits_per_text)
+
+ if not return_dict:
+ output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+ return ((loss,) + output) if loss is not None else output
+
+ return CLIPOutput(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+ @classmethod
+ def from_vision_text_pretrained(
+ cls,
+ vision_model_name_or_path: Optional[str] = None,
+ text_model_name_or_path: Optional[str] = None,
+ *model_args,
+ **kwargs,
+ ) -> PreTrainedModel:
+ """
+ Params:
+ vision_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the vision model. 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 *PyTorch checkpoint folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument. This
+ loading path is slower than converting the PyTorch checkpoint in a Flax model using the provided
+ conversion scripts and loading the Flax model afterwards.
+
+ text_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text model. 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 *PyTorch checkpoint folder* (e.g, `./pt_model`). In this case, `from_pt`
+ should be set to `True` and a configuration object should be provided as `config` argument. This
+ loading path is slower than converting the PyTorch checkpoint in a Flax model using the provided
+ conversion scripts and loading the Flax model afterwards.
+
+ 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 text configuration, use the prefix *text_* for each configuration parameter.
+ - To update the vision configuration, use the prefix *vision_* 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 VisionTextDualEncoderModel
+
+ >>> # initialize a model from pretrained ViT and BERT models. Note that the projection layers will be randomly initialized.
+ >>> model = VisionTextDualEncoderModel.from_vision_text_pretrained(
+ ... "google/vit-base-patch16-224", "google-bert/bert-base-uncased"
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./vit-bert")
+ >>> # load fine-tuned model
+ >>> model = VisionTextDualEncoderModel.from_pretrained("./vit-bert")
+ ```"""
+ kwargs_vision = {
+ argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
+ }
+
+ kwargs_text = {
+ argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
+ }
+
+ # remove vision, text kwargs from kwargs
+ for key in kwargs_vision.keys():
+ del kwargs["vision_" + key]
+ for key in kwargs_text.keys():
+ del kwargs["text_" + key]
+
+ # Load and initialize the vision and text model
+ vision_model = kwargs_vision.pop("model", None)
+ if vision_model is None:
+ if vision_model_name_or_path is None:
+ raise ValueError(
+ "If `vision_model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
+ )
+
+ if "config" not in kwargs_vision:
+ vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)
+
+ if vision_config.model_type == "clip":
+ kwargs_vision["config"] = vision_config.vision_config
+ vision_model = CLIPVisionModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)
+ # TODO: Should we use the pre-trained projection as well ?
+ else:
+ kwargs_vision["config"] = vision_config
+ vision_model = AutoModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)
+
+ text_model = kwargs_text.pop("model", None)
+ if text_model is None:
+ if text_model_name_or_path is None:
+ raise ValueError(
+ "If `text_model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
+ )
+
+ if "config" not in kwargs_text:
+ text_config = AutoConfig.from_pretrained(text_model_name_or_path)
+ kwargs_text["config"] = text_config
+
+ text_model = AutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)
+
+ # instantiate config with corresponding kwargs
+ config = VisionTextDualEncoderConfig.from_vision_text_configs(vision_model.config, text_model.config, **kwargs)
+
+ # init model
+ model = cls(config=config, vision_model=vision_model, text_model=text_model)
+
+ # the projection layers are always newly initialized when loading the model
+ # using pre-trained vision and text model.
+ logger.warning(
+ "The projection layer and logit scale weights `['visual_projection.weight', 'text_projection.weight',"
+ " 'logit_scale']` are newly initialized. You should probably TRAIN this model on a down-stream task to be"
+ " able to use it for predictions and inference."
+ )
+
+ return model
+
+
+__all__ = ["VisionTextDualEncoderModel"]
diff --git a/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/processing_vision_text_dual_encoder.py b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/processing_vision_text_dual_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..0b12bae8f7ef2aea1002048719a0883edbc76aad
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vision_text_dual_encoder/processing_vision_text_dual_encoder.py
@@ -0,0 +1,153 @@
+# 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.
+"""
+Processor class for VisionTextDualEncoder
+"""
+
+import warnings
+
+from ...processing_utils import ProcessorMixin
+from ...tokenization_utils_base import BatchEncoding
+
+
+class VisionTextDualEncoderProcessor(ProcessorMixin):
+ r"""
+ Constructs a VisionTextDualEncoder processor which wraps an image processor and a tokenizer into a single
+ processor.
+
+ [`VisionTextDualEncoderProcessor`] offers all the functionalities of [`AutoImageProcessor`] and [`AutoTokenizer`].
+ See the [`~VisionTextDualEncoderProcessor.__call__`] and [`~VisionTextDualEncoderProcessor.decode`] for more
+ information.
+
+ Args:
+ image_processor ([`AutoImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`PreTrainedTokenizer`], *optional*):
+ 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 have to specify an image_processor.")
+ if tokenizer is None:
+ raise ValueError("You have to specify a 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 VisionTextDualEncoderTokenizer's [`~PreTrainedTokenizer.__call__`] if `text` is not
+ `None` to encode the text. To prepare the image(s), this method forwards the `images` and `kwargs` arguments to
+ AutoImageProcessor's [`~AutoImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
+ 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 VisionTextDualEncoderTokenizer'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 VisionTextDualEncoderTokenizer'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
+
+
+__all__ = ["VisionTextDualEncoderProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/__init__.py b/lib/python3.12/site-packages/transformers/models/vivit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..13c52cbcff0b3fc6dc3255f41729754991dd3c8f
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vivit/__init__.py
@@ -0,0 +1,28 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_vivit import *
+ from .image_processing_vivit import *
+ from .modeling_vivit import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..142f7a7b9b4e8cb9a6500143ed16f3a1665187c7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/configuration_vivit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/configuration_vivit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..22b87c95be53d077327a33cf6976129a6175dfac
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/configuration_vivit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/image_processing_vivit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/image_processing_vivit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..588ab87d307dcb3057cae934ab41348df0e7ada2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/image_processing_vivit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/modeling_vivit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/modeling_vivit.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c3c71c19bdfb937fcd2451b17d9b696123d3e4c2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/vivit/__pycache__/modeling_vivit.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/configuration_vivit.py b/lib/python3.12/site-packages/transformers/models/vivit/configuration_vivit.py
new file mode 100644
index 0000000000000000000000000000000000000000..42863454e81c278ee7af62fee06b19ec9b6ca2f0
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vivit/configuration_vivit.py
@@ -0,0 +1,119 @@
+# 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.
+"""ViViT model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class VivitConfig(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`VivitModel`]. It is used to instantiate a ViViT
+ 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 ViViT
+ [google/vivit-b-16x2-kinetics400](https://huggingface.co/google/vivit-b-16x2-kinetics400) 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_frames (`int`, *optional*, defaults to 32):
+ The number of frames in each video.
+ tubelet_size (`List[int]`, *optional*, defaults to `[2, 16, 16]`):
+ The size (resolution) of each tubelet.
+ num_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ 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_fast"`):
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+ `"relu"`, `"selu"`, `"gelu_fast"` 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.
+ qkv_bias (`bool`, *optional*, defaults to `True`):
+ Whether to add a bias to the queries, keys and values.
+
+ Example:
+
+ ```python
+ >>> from transformers import VivitConfig, VivitModel
+
+ >>> # Initializing a ViViT google/vivit-b-16x2-kinetics400 style configuration
+ >>> configuration = VivitConfig()
+
+ >>> # Initializing a model (with random weights) from the google/vivit-b-16x2-kinetics400 style configuration
+ >>> model = VivitModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "vivit"
+
+ def __init__(
+ self,
+ image_size=224,
+ num_frames=32,
+ tubelet_size=[2, 16, 16],
+ num_channels=3,
+ hidden_size=768,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ intermediate_size=3072,
+ hidden_act="gelu_fast",
+ hidden_dropout_prob=0.0,
+ attention_probs_dropout_prob=0.0,
+ initializer_range=0.02,
+ layer_norm_eps=1e-06,
+ qkv_bias=True,
+ **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.num_frames = num_frames
+ self.tubelet_size = tubelet_size
+ self.num_channels = num_channels
+ self.qkv_bias = qkv_bias
+
+ super().__init__(**kwargs)
+
+
+__all__ = ["VivitConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/image_processing_vivit.py b/lib/python3.12/site-packages/transformers/models/vivit/image_processing_vivit.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b369be41ba883b0d9a853a043a7a17417024fc4
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vivit/image_processing_vivit.py
@@ -0,0 +1,407 @@
+# 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 processor class for Vivit."""
+
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+from transformers.utils import is_vision_available
+from transformers.utils.generic import TensorType
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from ...image_transforms import (
+ get_resize_output_image_size,
+ 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,
+ is_valid_image,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...utils import filter_out_non_signature_kwargs, logging
+
+
+if is_vision_available():
+ import PIL
+
+logger = logging.get_logger(__name__)
+
+
+def make_batched(videos) -> List[List[ImageInput]]:
+ if isinstance(videos, (list, tuple)) and isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0]):
+ return videos
+
+ elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]):
+ return [videos]
+
+ elif is_valid_image(videos):
+ return [[videos]]
+
+ raise ValueError(f"Could not make batched video from {videos}")
+
+
+class VivitImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a Vivit 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 the `preprocess` method.
+ size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 256}`):
+ Size of the output image after resizing. The shortest edge of the image will be resized to
+ `size["shortest_edge"]` while maintaining the aspect ratio of the original image. Can be overriden 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 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 the `do_center_crop`
+ parameter in the `preprocess` method.
+ crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Size of the image after applying the center crop. Can be overridden by the `crop_size` 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.
+ rescale_factor (`int` or `float`, *optional*, defaults to `1/127.5`):
+ Defines the scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter
+ in the `preprocess` method.
+ offset (`bool`, *optional*, defaults to `True`):
+ Whether to scale the image in both negative and positive directions. Can be overriden by the `offset` 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 = PILImageResampling.BILINEAR,
+ do_center_crop: bool = True,
+ crop_size: Dict[str, int] = None,
+ do_rescale: bool = True,
+ rescale_factor: Union[int, float] = 1 / 127.5,
+ offset: 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 {"shortest_edge": 256}
+ size = get_size_dict(size, default_to_square=False)
+ 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.do_center_crop = do_center_crop
+ self.crop_size = crop_size
+ self.resample = resample
+ self.do_rescale = do_rescale
+ self.rescale_factor = rescale_factor
+ self.offset = 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
+
+ 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.
+
+ Args:
+ image (`np.ndarray`):
+ Image to resize.
+ size (`Dict[str, int]`):
+ Size of the output image. If `size` is of the form `{"height": h, "width": w}`, the output image will
+ have the size `(h, w)`. If `size` is of the form `{"shortest_edge": s}`, the output image will have its
+ shortest edge of length `s` while keeping the aspect ratio of the original image.
+ resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
+ 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, default_to_square=False)
+ if "shortest_edge" in size:
+ output_size = get_resize_output_image_size(
+ image, size["shortest_edge"], default_to_square=False, input_data_format=input_data_format
+ )
+ elif "height" in size and "width" in size:
+ output_size = (size["height"], size["width"])
+ else:
+ raise ValueError(f"Size must have 'height' and 'width' or 'shortest_edge' as keys. Got {size.keys()}")
+ return resize(
+ image,
+ size=output_size,
+ resample=resample,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ **kwargs,
+ )
+
+ # Copied from transformers.models.efficientnet.image_processing_efficientnet.EfficientNetImageProcessor.rescale
+ 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_image(
+ self,
+ image: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: Dict[str, int] = None,
+ do_rescale: Optional[bool] = None,
+ rescale_factor: Optional[float] = None,
+ offset: Optional[bool] = None,
+ do_normalize: Optional[bool] = None,
+ image_mean: Optional[Union[float, List[float]]] = None,
+ image_std: Optional[Union[float, List[float]]] = None,
+ data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, 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,
+ )
+
+ if offset and not do_rescale:
+ raise ValueError("For offset, do_rescale must also be set to True.")
+
+ # All transformations expect numpy arrays.
+ image = to_numpy_array(image)
+
+ if do_rescale and is_scaled_image(image):
+ 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:
+ 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, size=crop_size, input_data_format=input_data_format)
+
+ if do_rescale:
+ image = self.rescale(image=image, scale=rescale_factor, offset=offset, 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)
+
+ image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+ return image
+
+ @filter_out_non_signature_kwargs()
+ def preprocess(
+ self,
+ videos: ImageInput,
+ do_resize: Optional[bool] = None,
+ size: Dict[str, int] = None,
+ resample: PILImageResampling = None,
+ do_center_crop: Optional[bool] = None,
+ crop_size: Dict[str, int] = None,
+ do_rescale: Optional[bool] = None,
+ rescale_factor: Optional[float] = None,
+ offset: Optional[bool] = 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: ChannelDimension = ChannelDimension.FIRST,
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
+ ) -> PIL.Image.Image:
+ """
+ Preprocess an image or batch of images.
+
+ Args:
+ videos (`ImageInput`):
+ Video frames to preprocess. Expects a single or batch of video frames with pixel values ranging from 0
+ to 255. If passing in frames 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 applying resize.
+ resample (`PILImageResampling`, *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_centre_crop`):
+ Whether to centre crop the image.
+ crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
+ Size of the image after applying the centre crop.
+ do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+ Whether to rescale the image values between `[-1 - 1]` if `offset` is `True`, `[0, 1]` otherwise.
+ rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+ Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+ offset (`bool`, *optional*, defaults to `self.offset`):
+ Whether to scale the image in both negative and positive directions.
+ 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.
+ - Unset: Use the inferred 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
+ 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
+ offset = offset if offset is not None else self.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
+
+ size = size if size is not None else self.size
+ size = get_size_dict(size, default_to_square=False)
+ 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")
+
+ if not valid_images(videos):
+ raise ValueError(
+ "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+ "torch.Tensor, tf.Tensor or jax.ndarray."
+ )
+
+ videos = make_batched(videos)
+
+ videos = [
+ [
+ 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,
+ offset=offset,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+ for img in video
+ ]
+ for video in videos
+ ]
+
+ data = {"pixel_values": videos}
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["VivitImageProcessor"]
diff --git a/lib/python3.12/site-packages/transformers/models/vivit/modeling_vivit.py b/lib/python3.12/site-packages/transformers/models/vivit/modeling_vivit.py
new file mode 100644
index 0000000000000000000000000000000000000000..669106239a06a83cb1f4737f76eb7a5bed474227
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/vivit/modeling_vivit.py
@@ -0,0 +1,844 @@
+# coding=utf-8
+# Copyright 2023 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.
+"""PyTorch ViViT model."""
+
+from typing import Callable, Optional, Set, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
+from ...utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+ torch_int,
+)
+from .configuration_vivit import VivitConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "google/vivit-b-16x2-kinetics400"
+_CONFIG_FOR_DOC = "VivitConfig"
+
+
+class VivitTubeletEmbeddings(nn.Module):
+ """
+ Construct Vivit Tubelet embeddings.
+
+ This module turns a batch of videos of shape (batch_size, num_frames, num_channels, height, width) into a tensor of
+ shape (batch_size, seq_len, hidden_size) to be consumed by a Transformer encoder.
+
+ The seq_len (the number of patches) equals (number of frames // tubelet_size[0]) * (height // tubelet_size[1]) *
+ (width // tubelet_size[2]).
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.num_frames = config.num_frames
+ self.image_size = config.image_size
+ self.patch_size = config.tubelet_size
+ self.num_patches = (
+ (self.image_size // self.patch_size[2])
+ * (self.image_size // self.patch_size[1])
+ * (self.num_frames // self.patch_size[0])
+ )
+ self.embed_dim = config.hidden_size
+
+ self.projection = nn.Conv3d(
+ config.num_channels, config.hidden_size, kernel_size=config.tubelet_size, stride=config.tubelet_size
+ )
+
+ def forward(self, pixel_values, interpolate_pos_encoding: bool = False):
+ batch_size, num_frames, num_channels, height, width = pixel_values.shape
+ if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size):
+ raise ValueError(
+ f"Image image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."
+ )
+
+ # permute to (batch_size, num_channels, num_frames, height, width)
+ pixel_values = pixel_values.permute(0, 2, 1, 3, 4)
+
+ x = self.projection(pixel_values)
+ # out_batch_size, out_num_channels, out_num_frames, out_height, out_width = x.shape
+ # flattens time and space dimensions, transposes to (out_batch_size, flat_tokens, out_num_channels)
+ x = x.flatten(2).transpose(1, 2)
+ return x
+
+
+class VivitEmbeddings(nn.Module):
+ """
+ Vivit Embeddings.
+
+ Creates embeddings from a video using VivitTubeletEmbeddings, adds CLS token and positional embeddings.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+ self.patch_embeddings = VivitTubeletEmbeddings(config)
+
+ self.position_embeddings = nn.Parameter(
+ torch.zeros(1, self.patch_embeddings.num_patches + 1, config.hidden_size)
+ )
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.patch_size = config.tubelet_size[1:]
+ self.config = config
+
+ # Adapted from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding
+ 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. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embeddings.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, :1]
+ patch_pos_embed = self.position_embeddings[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size[0]
+ new_width = width // self.patch_size[1]
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values, interpolate_pos_encoding: bool = False):
+ batch_size, num_frames, num_channels, height, width = pixel_values.shape
+ embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+
+ cls_tokens = self.cls_token.tile([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
+
+
+# Copied from transformers.models.vit.modeling_vit.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+
+ # Normalize the attention scores to probabilities.
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ # Mask heads if we want to
+ if attention_mask is not None:
+ attn_weights = attn_weights * attention_mask
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Vivit
+class VivitSelfAttention(nn.Module):
+ def __init__(self, config: VivitConfig) -> 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.config = config
+ 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.dropout_prob = config.attention_probs_dropout_prob
+ self.scaling = self.attention_head_size**-0.5
+ self.is_causal = False
+
+ 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)
+
+ 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]]:
+ 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(self.query(hidden_states))
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ if self.config._attn_implementation == "sdpa" and output_attentions:
+ logger.warning_once(
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+ )
+ else:
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ context_layer, attention_probs = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ head_mask,
+ is_causal=self.is_causal,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.dropout_prob,
+ )
+
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.reshape(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->Vivit
+class VivitSelfOutput(nn.Module):
+ """
+ The residual connection is defined in VivitLayer instead of here (as is the case with other models), due to the
+ layernorm applied before each block.
+ """
+
+ def __init__(self, config: VivitConfig) -> 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->Vivit
+class VivitAttention(nn.Module):
+ def __init__(self, config: VivitConfig) -> None:
+ super().__init__()
+ self.attention = VivitSelfAttention(config)
+ self.output = VivitSelfOutput(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 VivitIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ 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)
+ hidden_states = self.dropout(hidden_states)
+
+ return hidden_states
+
+
+class VivitOutput(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 VivitLayer(nn.Module):
+ """This corresponds to the EncoderBlock class in the scenic/vivit implementation."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = VivitAttention(config)
+ self.intermediate = VivitIntermediate(config)
+ self.output = VivitOutput(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, head_mask=None, output_attentions=False):
+ self_attention_outputs = self.attention(
+ # in Vivit, layernorm is applied before self-attention
+ self.layernorm_before(hidden_states),
+ head_mask,
+ output_attentions=output_attentions,
+ )
+ attention_output = self_attention_outputs[0]
+ # add self attentions if we output attention weights
+ outputs = self_attention_outputs[1:]
+
+ # first residual connection
+ hidden_states = attention_output + hidden_states
+
+ # in Vivit, 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 VivitEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([VivitLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states,
+ head_mask=None,
+ 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
+
+ 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 VivitPooler(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):
+ # 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 VivitPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = VivitConfig
+ base_model_prefix = "vivit"
+ main_input_name = "pixel_values"
+ supports_gradient_checkpointing = True
+ _no_split_modules = []
+ _supports_sdpa = True
+ _supports_flash_attn_2 = True
+
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ if isinstance(module, (nn.Linear, nn.Conv3d)):
+ # 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, VivitEmbeddings):
+ module.cls_token.data.zero_()
+ module.position_embeddings.data.zero_()
+
+
+VIVIT_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 ([`VivitConfig`]): 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.
+"""
+
+VIVIT_INPUTS_DOCSTRING = r"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+ Pixel values. Pixel values can be obtained using [`VivitImageProcessor`]. See
+ [`VivitImageProcessor.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.
+ interpolate_pos_encoding (`bool`, *optional*, `False`):
+ Whether to interpolate the pre-trained position encodings.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare ViViT Transformer model outputting raw hidden-states without any specific head on top.",
+ VIVIT_START_DOCSTRING,
+)
+class VivitModel(VivitPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=True):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = VivitEmbeddings(config)
+ self.encoder = VivitEncoder(config)
+
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.pooler = VivitPooler(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.
+
+ Args:
+ heads_to_prune:
+ dict of {layer_num: list of heads to prune in this layer}
+ """
+ for layer, heads in heads_to_prune.items():
+ self.encoder.layer[layer].attention.prune_heads(heads)
+
+ @add_start_docstrings_to_model_forward(VIVIT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ head_mask: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple[torch.FloatTensor], BaseModelOutputWithPooling]:
+ r"""
+ Returns:
+
+ Examples:
+
+ ```python
+ >>> import av
+ >>> import numpy as np
+
+ >>> from transformers import VivitImageProcessor, VivitModel
+ >>> from huggingface_hub import hf_hub_download
+
+ >>> np.random.seed(0)
+
+
+ >>> 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
+
+
+ >>> # 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"
+ ... )
+ >>> container = av.open(file_path)
+
+ >>> # sample 32 frames
+ >>> indices = sample_frame_indices(clip_len=32, frame_sample_rate=1, seg_len=container.streams.video[0].frames)
+ >>> video = read_video_pyav(container=container, indices=indices)
+
+ >>> image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
+ >>> model = VivitModel.from_pretrained("google/vivit-b-16x2-kinetics400")
+
+ >>> # prepare video for the model
+ >>> inputs = image_processor(list(video), return_tensors="pt")
+
+ >>> # forward pass
+ >>> outputs = model(**inputs)
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 3137, 768]
+ ```"""
+ 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")
+
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+ embedding_output = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+
+ 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:
+ 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(
+ """
+ ViViT Transformer model with a video classification head on top (a linear layer on top of the final hidden state of the
+[CLS] token) e.g. for Kinetics-400.
+
+
+
+ Note that it's possible to fine-tune ViT on higher resolution images than the ones it has been trained on, by
+ setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained
+ position embeddings to the higher resolution.
+
+
+ """,
+ VIVIT_START_DOCSTRING,
+)
+class VivitForVideoClassification(VivitPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.vivit = VivitModel(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(VIVIT_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=ImageClassifierOutput, 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,
+ interpolate_pos_encoding: bool = False,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple[torch.FloatTensor], 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
+ >>> import av
+ >>> import numpy as np
+ >>> import torch
+
+ >>> from transformers import VivitImageProcessor, VivitForVideoClassification
+ >>> from huggingface_hub import hf_hub_download
+
+ >>> np.random.seed(0)
+
+
+ >>> 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
+
+
+ >>> # 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"
+ ... )
+ >>> container = av.open(file_path)
+
+ >>> # sample 32 frames
+ >>> indices = sample_frame_indices(clip_len=32, frame_sample_rate=4, seg_len=container.streams.video[0].frames)
+ >>> video = read_video_pyav(container=container, indices=indices)
+
+ >>> image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
+ >>> model = VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400")
+
+ >>> inputs = image_processor(list(video), return_tensors="pt")
+
+ >>> with torch.no_grad():
+ ... outputs = model(**inputs)
+ ... logits = outputs.logits
+
+ >>> # model predicts one of the 400 Kinetics-400 classes
+ >>> predicted_label = logits.argmax(-1).item()
+ >>> print(model.config.id2label[predicted_label])
+ LABEL_116
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.vivit(
+ pixel_values,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.classifier(sequence_output[:, 0, :])
+
+ 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 ImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["VivitModel", "VivitPreTrainedModel", "VivitForVideoClassification"]
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__init__.py b/lib/python3.12/site-packages/transformers/models/xlnet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f453455925328db06bf86ecc85c5acc8762dd26
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/__init__.py
@@ -0,0 +1,30 @@
+# 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 _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_xlnet import *
+ from .modeling_tf_xlnet import *
+ from .modeling_xlnet import *
+ from .tokenization_xlnet import *
+ from .tokenization_xlnet_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..492685791d280a30318d6adffc7447278edc700e
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/configuration_xlnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/configuration_xlnet.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ddc409003832f0960036b61d7ef78e2af253984a
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/configuration_xlnet.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_tf_xlnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_tf_xlnet.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e3e01508bdb10fc12845b58537efd3992527539a
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_tf_xlnet.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_xlnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_xlnet.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e7363bd318a050b9c307863129dd193483d53ced
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/modeling_xlnet.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c9b93caa33a7a7ea9a8c29b2aaae3d2197457152
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet_fast.cpython-312.pyc b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet_fast.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..14c437e591a7b84e235b7493f63fb4d21a67ba28
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/models/xlnet/__pycache__/tokenization_xlnet_fast.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/configuration_xlnet.py b/lib/python3.12/site-packages/transformers/models/xlnet/configuration_xlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a7238eb4c098b93c7aa911bb958ebe771a4539c
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/configuration_xlnet.py
@@ -0,0 +1,240 @@
+# 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.
+"""XLNet configuration"""
+
+import warnings
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class XLNetConfig(PretrainedConfig):
+ """
+ This is the configuration class to store the configuration of a [`XLNetModel`] or a [`TFXLNetModel`]. It is used to
+ instantiate a XLNet 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
+ [xlnet/xlnet-large-cased](https://huggingface.co/xlnet/xlnet-large-cased) 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 XLNet model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`XLNetModel`] or [`TFXLNetModel`].
+ d_model (`int`, *optional*, defaults to 1024):
+ Dimensionality of the encoder layers and the pooler layer.
+ n_layer (`int`, *optional*, defaults to 24):
+ 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.
+ d_inner (`int`, *optional*, defaults to 4096):
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+ ff_activation (`str` or `Callable`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the If string, `"gelu"`, `"relu"`, `"silu"` and
+ `"gelu_new"` are supported.
+ untie_r (`bool`, *optional*, defaults to `True`):
+ Whether or not to untie relative position biases
+ attn_type (`str`, *optional*, defaults to `"bi"`):
+ The attention type used by the model. Set `"bi"` for XLNet, `"uni"` for Transformer-XL.
+ 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.
+ dropout (`float`, *optional*, defaults to 0.1):
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+ mem_len (`int` or `None`, *optional*):
+ The number of tokens to cache. The key/value pairs that have already been pre-computed in a previous
+ forward pass won't be re-computed. See the
+ [quickstart](https://huggingface.co/transformers/quickstart.html#using-the-past) for more information.
+ reuse_len (`int`, *optional*):
+ The number of tokens in the current batch to be cached and reused in the future.
+ bi_data (`bool`, *optional*, defaults to `False`):
+ Whether or not to use bidirectional input pipeline. Usually set to `True` during pretraining and `False`
+ during finetuning.
+ clamp_len (`int`, *optional*, defaults to -1):
+ Clamp all relative distances larger than clamp_len. Setting this attribute to -1 means no clamping.
+ same_length (`bool`, *optional*, defaults to `False`):
+ Whether or not to use the same attention length for each token.
+ summary_type (`str`, *optional*, defaults to "last"):
+ 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 (`boo`, *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_last_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.
+ use_mems_eval (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should make use of the recurrent memory mechanism in evaluation mode.
+ use_mems_train (`bool`, *optional*, defaults to `False`):
+ Whether or not the model should make use of the recurrent memory mechanism in train mode.
+
+
+
+ For pretraining, it is recommended to set `use_mems_train` to `True`. For fine-tuning, it is recommended to
+ set `use_mems_train` to `False` as discussed
+ [here](https://github.com/zihangdai/xlnet/issues/41#issuecomment-505102587). If `use_mems_train` is set to
+ `True`, one has to make sure that the train batches are correctly pre-processed, *e.g.* `batch_1 = [[This
+ line is], [This is the]]` and `batch_2 = [[ the first line], [ second line]]` and that all batches are of
+ equal size.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import XLNetConfig, XLNetModel
+
+ >>> # Initializing a XLNet configuration
+ >>> configuration = XLNetConfig()
+
+ >>> # Initializing a model (with random weights) from the configuration
+ >>> model = XLNetModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "xlnet"
+ keys_to_ignore_at_inference = ["mems"]
+ attribute_map = {
+ "n_token": "vocab_size", # Backward compatibility
+ "hidden_size": "d_model",
+ "num_attention_heads": "n_head",
+ "num_hidden_layers": "n_layer",
+ }
+
+ def __init__(
+ self,
+ vocab_size=32000,
+ d_model=1024,
+ n_layer=24,
+ n_head=16,
+ d_inner=4096,
+ ff_activation="gelu",
+ untie_r=True,
+ attn_type="bi",
+ initializer_range=0.02,
+ layer_norm_eps=1e-12,
+ dropout=0.1,
+ mem_len=512,
+ reuse_len=None,
+ use_mems_eval=True,
+ use_mems_train=False,
+ bi_data=False,
+ clamp_len=-1,
+ same_length=False,
+ summary_type="last",
+ summary_use_proj=True,
+ summary_activation="tanh",
+ summary_last_dropout=0.1,
+ start_n_top=5,
+ end_n_top=5,
+ pad_token_id=5,
+ bos_token_id=1,
+ eos_token_id=2,
+ **kwargs,
+ ):
+ """Constructs XLNetConfig."""
+ self.vocab_size = vocab_size
+ self.d_model = d_model
+ self.n_layer = n_layer
+ self.n_head = n_head
+ if d_model % n_head != 0:
+ raise ValueError(f"'d_model % n_head' ({d_model % n_head}) should be equal to 0")
+ if "d_head" in kwargs:
+ if kwargs["d_head"] != d_model // n_head:
+ raise ValueError(
+ f"`d_head` ({kwargs['d_head']}) should be equal to `d_model // n_head` ({d_model // n_head})"
+ )
+ self.d_head = d_model // n_head
+ self.ff_activation = ff_activation
+ self.d_inner = d_inner
+ self.untie_r = untie_r
+ self.attn_type = attn_type
+
+ self.initializer_range = initializer_range
+ self.layer_norm_eps = layer_norm_eps
+
+ self.dropout = dropout
+ self.mem_len = mem_len
+ self.reuse_len = reuse_len
+ self.bi_data = bi_data
+ self.clamp_len = clamp_len
+ self.same_length = same_length
+
+ 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.start_n_top = start_n_top
+ self.end_n_top = end_n_top
+
+ self.bos_token_id = bos_token_id
+ self.pad_token_id = pad_token_id
+ self.eos_token_id = eos_token_id
+
+ if "use_cache" in kwargs:
+ warnings.warn(
+ "The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`"
+ " instead.",
+ FutureWarning,
+ )
+ use_mems_eval = kwargs["use_cache"]
+
+ self.use_mems_eval = use_mems_eval
+ self.use_mems_train = use_mems_train
+ super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+ @property
+ def max_position_embeddings(self):
+ 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."
+ )
+
+
+__all__ = ["XLNetConfig"]
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/modeling_tf_xlnet.py b/lib/python3.12/site-packages/transformers/models/xlnet/modeling_tf_xlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..2383352ae6af0cafb538aef15e8c8c66c3501d43
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/modeling_tf_xlnet.py
@@ -0,0 +1,1821 @@
+# 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 XLNet model.
+"""
+
+from __future__ import annotations
+
+import warnings
+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_utils import (
+ TFCausalLanguageModelingLoss,
+ 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 (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from .configuration_xlnet import XLNetConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "xlnet/xlnet-base-cased"
+_CONFIG_FOR_DOC = "XLNetConfig"
+
+
+class TFXLNetRelativeAttention(keras.layers.Layer):
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+
+ if config.d_model % config.n_head != 0:
+ raise ValueError(
+ f"The hidden size ({config.d_model}) is not a multiple of the number of attention "
+ f"heads ({config.n_head}"
+ )
+
+ self.n_head = config.n_head
+ self.d_head = config.d_head
+ self.d_model = config.d_model
+ self.scale = 1 / (config.d_head**0.5)
+ self.initializer_range = config.initializer_range
+ self.output_attentions = config.output_attentions
+
+ self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
+ self.dropout = keras.layers.Dropout(config.dropout)
+ self.config = config
+
+ def build(self, input_shape=None):
+ initializer = get_initializer(self.initializer_range)
+ self.q = self.add_weight(
+ shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="q"
+ )
+ self.k = self.add_weight(
+ shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="k"
+ )
+ self.v = self.add_weight(
+ shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="v"
+ )
+ self.o = self.add_weight(
+ shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="o"
+ )
+ self.r = self.add_weight(
+ shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="r"
+ )
+ self.r_r_bias = self.add_weight(
+ shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_r_bias"
+ )
+ self.r_s_bias = self.add_weight(
+ shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_s_bias"
+ )
+ self.r_w_bias = self.add_weight(
+ shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_w_bias"
+ )
+ self.seg_embed = self.add_weight(
+ shape=(2, self.n_head, self.d_head), initializer=initializer, trainable=True, name="seg_embed"
+ )
+
+ if self.built:
+ return
+ self.built = True
+ 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])
+
+ def prune_heads(self, heads):
+ raise NotImplementedError
+
+ def rel_shift(self, x, klen=-1):
+ """perform relative shift to form the relative attention score."""
+ x_size = shape_list(x)
+
+ x = tf.reshape(x, (x_size[1], x_size[0], x_size[2], x_size[3]))
+ x = x[1:, ...]
+ x = tf.reshape(x, (x_size[0], x_size[1] - 1, x_size[2], x_size[3]))
+ x = x[:, 0:klen, :, :]
+ # x = torch.index_select(x, 1, torch.arange(klen, device=x.device, dtype=torch.long))
+
+ return x
+
+ def rel_attn_core(
+ self, q_head, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask, head_mask, output_attentions, training=False
+ ):
+ """Core relative positional attention operations."""
+ # content based attention score
+ ac = tf.einsum("ibnd,jbnd->ijbn", q_head + self.r_w_bias, k_head_h)
+
+ # position based attention score
+ bd = tf.einsum("ibnd,jbnd->ijbn", q_head + self.r_r_bias, k_head_r)
+ bd = self.rel_shift(bd, klen=shape_list(ac)[1])
+
+ # segment based attention score
+ if seg_mat is None:
+ ef = 0
+ else:
+ ef = tf.einsum("ibnd,snd->ibns", q_head + self.r_s_bias, self.seg_embed)
+ ef = tf.einsum("ijbs,ibns->ijbn", seg_mat, ef)
+
+ # merge attention scores and perform masking
+ attn_score = (ac + bd + ef) * self.scale
+ if attn_mask is not None:
+ # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
+ if attn_mask.dtype == tf.float16 or attn_mask.dtype == tf.bfloat16:
+ attn_score = attn_score - 65500 * attn_mask
+ else:
+ attn_score = attn_score - 1e30 * attn_mask
+
+ # attention probability
+ attn_prob = stable_softmax(attn_score, axis=1)
+
+ attn_prob = self.dropout(attn_prob, training=training)
+
+ # Mask heads if we want to
+ if head_mask is not None:
+ attn_prob = attn_prob * head_mask
+
+ # attention output
+ attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, v_head_h)
+
+ if output_attentions:
+ return attn_vec, attn_prob
+
+ return attn_vec
+
+ def post_attention(self, h, attn_vec, residual=True, training=False):
+ """Post-attention processing."""
+ # post-attention projection (back to `d_model`)
+ attn_out = tf.einsum("ibnd,hnd->ibh", attn_vec, self.o)
+
+ attn_out = self.dropout(attn_out, training=training)
+
+ if residual:
+ attn_out = attn_out + h
+ output = self.layer_norm(attn_out)
+
+ return output
+
+ def call(
+ self,
+ h,
+ g,
+ attn_mask_h,
+ attn_mask_g,
+ r,
+ seg_mat,
+ mems: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ output_attentions: Optional[bool] = False,
+ training: bool = False,
+ ):
+ if g is not None:
+ # Two-stream attention with relative positional encoding.
+ # content based attention score
+ if mems is not None and len(shape_list(mems)) > 1:
+ cat = tf.concat([mems, h], axis=0)
+ else:
+ cat = h
+
+ # content-based key head
+ k_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.k)
+
+ # content-based value head
+ v_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.v)
+
+ # position-based key head
+ k_head_r = tf.einsum("ibh,hnd->ibnd", r, self.r)
+
+ # h-stream
+ # content-stream query head
+ q_head_h = tf.einsum("ibh,hnd->ibnd", h, self.q)
+
+ # core attention ops
+ attn_vec_h = self.rel_attn_core(
+ q_head_h,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat,
+ attn_mask_h,
+ head_mask,
+ output_attentions,
+ training=training,
+ )
+
+ if output_attentions:
+ attn_vec_h, attn_prob_h = attn_vec_h
+
+ # post processing
+ output_h = self.post_attention(h, attn_vec_h, training=training)
+
+ # g-stream
+ # query-stream query head
+ q_head_g = tf.einsum("ibh,hnd->ibnd", g, self.q)
+
+ # core attention ops
+ if target_mapping is not None:
+ q_head_g = tf.einsum("mbnd,mlb->lbnd", q_head_g, target_mapping)
+ attn_vec_g = self.rel_attn_core(
+ q_head_g,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat,
+ attn_mask_g,
+ head_mask,
+ output_attentions,
+ training=training,
+ )
+
+ if output_attentions:
+ attn_vec_g, attn_prob_g = attn_vec_g
+
+ attn_vec_g = tf.einsum("lbnd,mlb->mbnd", attn_vec_g, target_mapping)
+ else:
+ attn_vec_g = self.rel_attn_core(
+ q_head_g,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat,
+ attn_mask_g,
+ head_mask,
+ output_attentions,
+ training=training,
+ )
+
+ if output_attentions:
+ attn_vec_g, attn_prob_g = attn_vec_g
+
+ # post processing
+ output_g = self.post_attention(g, attn_vec_g, training=training)
+
+ if output_attentions:
+ attn_prob = attn_prob_h, attn_prob_g
+
+ else:
+ # Multi-head attention with relative positional encoding
+ if mems is not None and len(shape_list(mems)) > 1:
+ cat = tf.concat([mems, h], axis=0)
+ else:
+ cat = h
+
+ # content heads
+ q_head_h = tf.einsum("ibh,hnd->ibnd", h, self.q)
+ k_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.k)
+ v_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.v)
+
+ # positional heads
+ k_head_r = tf.einsum("ibh,hnd->ibnd", r, self.r)
+
+ # core attention ops
+ attn_vec = self.rel_attn_core(
+ q_head_h,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat,
+ attn_mask_h,
+ head_mask,
+ output_attentions,
+ training=training,
+ )
+
+ if output_attentions:
+ attn_vec, attn_prob = attn_vec
+
+ # post processing
+ output_h = self.post_attention(h, attn_vec, training=training)
+ output_g = None
+
+ outputs = (output_h, output_g)
+ if output_attentions:
+ outputs = outputs + (attn_prob,)
+ return outputs
+
+
+class TFXLNetFeedForward(keras.layers.Layer):
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+ self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
+ self.layer_1 = keras.layers.Dense(
+ config.d_inner, kernel_initializer=get_initializer(config.initializer_range), name="layer_1"
+ )
+ self.layer_2 = keras.layers.Dense(
+ config.d_model, kernel_initializer=get_initializer(config.initializer_range), name="layer_2"
+ )
+ self.dropout = keras.layers.Dropout(config.dropout)
+ if isinstance(config.ff_activation, str):
+ self.activation_function = get_tf_activation(config.ff_activation)
+ else:
+ self.activation_function = config.ff_activation
+ self.config = config
+
+ def call(self, inp, training=False):
+ output = inp
+ output = self.layer_1(output)
+ output = self.activation_function(output)
+ output = self.dropout(output, training=training)
+ output = self.layer_2(output)
+ output = self.dropout(output, training=training)
+ output = self.layer_norm(output + inp)
+ return output
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ 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])
+ if getattr(self, "layer_1", None) is not None:
+ with tf.name_scope(self.layer_1.name):
+ self.layer_1.build([None, None, self.config.d_model])
+ if getattr(self, "layer_2", None) is not None:
+ with tf.name_scope(self.layer_2.name):
+ self.layer_2.build([None, None, self.config.d_inner])
+
+
+class TFXLNetLayer(keras.layers.Layer):
+ def __init__(self, config, **kwargs):
+ super().__init__(**kwargs)
+ self.rel_attn = TFXLNetRelativeAttention(config, name="rel_attn")
+ self.ff = TFXLNetFeedForward(config, name="ff")
+ self.dropout = keras.layers.Dropout(config.dropout)
+
+ def call(
+ self,
+ output_h,
+ output_g,
+ non_tgt_mask,
+ attn_mask,
+ pos_emb,
+ seg_mat,
+ mems: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ output_attentions: Optional[bool] = False,
+ training: bool = False,
+ ):
+ outputs = self.rel_attn(
+ output_h,
+ output_g,
+ non_tgt_mask,
+ attn_mask,
+ pos_emb,
+ seg_mat,
+ mems,
+ target_mapping,
+ head_mask,
+ output_attentions,
+ training=training,
+ )
+ output_h, output_g = outputs[:2]
+
+ if output_g is not None:
+ output_g = self.ff(output_g, training=training)
+ output_h = self.ff(output_h, training=training)
+
+ outputs = (output_h, output_g) + outputs[2:] # Add again attentions if there are there
+ return outputs
+
+ def build(self, input_shape=None):
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "rel_attn", None) is not None:
+ with tf.name_scope(self.rel_attn.name):
+ self.rel_attn.build(None)
+ if getattr(self, "ff", None) is not None:
+ with tf.name_scope(self.ff.name):
+ self.ff.build(None)
+
+
+class TFXLNetLMHead(keras.layers.Layer):
+ def __init__(self, config, input_embeddings, **kwargs):
+ super().__init__(**kwargs)
+ self.config = config
+ # 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):
+ hidden_states = self.input_embeddings(hidden_states, mode="linear")
+ hidden_states = hidden_states + self.bias
+ return hidden_states
+
+
+@keras_serializable
+class TFXLNetMainLayer(keras.layers.Layer):
+ config_class = XLNetConfig
+
+ 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.return_dict
+
+ self.mem_len = config.mem_len
+ self.reuse_len = config.reuse_len
+ self.d_model = config.d_model
+ self.same_length = config.same_length
+ self.attn_type = config.attn_type
+ self.bi_data = config.bi_data
+ self.clamp_len = config.clamp_len
+ self.n_layer = config.n_layer
+ self.use_bfloat16 = config.use_bfloat16
+ self.initializer_range = config.initializer_range
+
+ self.word_embedding = TFSharedEmbeddings(
+ config.vocab_size, config.d_model, initializer_range=config.initializer_range, name="word_embedding"
+ )
+ self.layer = [TFXLNetLayer(config, name=f"layer_._{i}") for i in range(config.n_layer)]
+ self.dropout = keras.layers.Dropout(config.dropout)
+
+ self.use_mems_eval = config.use_mems_eval
+ self.use_mems_train = config.use_mems_train
+
+ def get_input_embeddings(self):
+ return self.word_embedding
+
+ def set_input_embeddings(self, value):
+ self.word_embedding.weight = value
+ self.word_embedding.vocab_size = shape_list(value)[0]
+
+ def build(self, input_shape=None):
+ initializer = get_initializer(self.initializer_range)
+ self.mask_emb = self.add_weight(
+ shape=(1, 1, self.d_model), initializer=initializer, trainable=True, name="mask_emb"
+ )
+
+ if self.built:
+ return
+ self.built = True
+ if getattr(self, "word_embedding", None) is not None:
+ with tf.name_scope(self.word_embedding.name):
+ self.word_embedding.build(None)
+ if getattr(self, "layer", None) is not None:
+ for layer in self.layer:
+ with tf.name_scope(layer.name):
+ layer.build(None)
+
+ def _prune_heads(self, heads_to_prune):
+ raise NotImplementedError
+
+ def create_mask(self, qlen, mlen):
+ """
+ Creates causal attention mask. Float mask where 1.0 indicates masked, 0.0 indicates not-masked.
+
+ Args:
+ qlen: TODO Lysandre didn't fill
+ mlen: TODO Lysandre didn't fill
+
+ ```
+
+ same_length=False: same_length=True:
+ < qlen > < qlen >
+ ^ [0 0 0 0 0 1 1 1 1] [0 0 0 0 0 1 1 1 1]
+ [0 0 0 0 0 0 1 1 1] [1 0 0 0 0 0 1 1 1]
+ qlen [0 0 0 0 0 0 0 1 1] [1 1 0 0 0 0 0 1 1]
+ [0 0 0 0 0 0 0 0 1] [1 1 1 0 0 0 0 0 1]
+ v [0 0 0 0 0 0 0 0 0] [1 1 1 1 0 0 0 0 0]
+ ```
+ """
+ attn_mask = tf.ones([qlen, qlen])
+ mask_u = tf.linalg.band_part(attn_mask, 0, -1)
+ mask_dia = tf.linalg.band_part(attn_mask, 0, 0)
+ attn_mask_pad = tf.zeros([qlen, mlen])
+ ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
+ if self.same_length:
+ mask_l = tf.linalg.band_part(attn_mask, -1, 0)
+ ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)
+ return ret
+
+ def cache_mem(self, curr_out, prev_mem):
+ # cache hidden states into memory.
+ if self.reuse_len is not None and self.reuse_len > 0:
+ curr_out = curr_out[: self.reuse_len]
+
+ if self.mem_len is None or self.mem_len == 0:
+ # If `use_mems` is active but no `mem_len` is defined, the model behaves like GPT-2 at inference time
+ # and returns all of the past and current hidden states.
+ cutoff = 0
+ else:
+ # If `use_mems` is active and `mem_len` is defined, the model returns the last `mem_len` hidden
+ # states. This is the preferred setting for training and long-form generation.
+ cutoff = -self.mem_len
+ if prev_mem is None:
+ # if `use_mems` is active and `mem_len` is defined, the model
+ new_mem = curr_out[cutoff:]
+ else:
+ new_mem = tf.concat([prev_mem, curr_out], 0)[cutoff:]
+
+ return tf.stop_gradient(new_mem)
+
+ @staticmethod
+ def positional_embedding(pos_seq, inv_freq, bsz=None):
+ sinusoid_inp = tf.einsum("i,d->id", pos_seq, inv_freq)
+ pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], axis=-1)
+ pos_emb = pos_emb[:, None, :]
+
+ if bsz is not None:
+ pos_emb = tf.tile(pos_emb, [1, bsz, 1])
+
+ return pos_emb
+
+ def relative_positional_encoding(self, qlen, klen, bsz=None):
+ """create relative positional encoding."""
+ freq_seq = tf.range(0, self.d_model, 2.0)
+ inv_freq = 1 / (10000 ** (freq_seq / self.d_model))
+
+ if self.attn_type == "bi":
+ # beg, end = klen - 1, -qlen
+ beg, end = klen, -qlen
+ elif self.attn_type == "uni":
+ # beg, end = klen - 1, -1
+ beg, end = klen, -1
+ else:
+ raise ValueError(f"Unknown `attn_type` {self.attn_type}.")
+
+ if self.bi_data:
+ fwd_pos_seq = tf.range(beg, end, -1.0)
+ bwd_pos_seq = tf.range(-beg, -end, 1.0)
+
+ if self.clamp_len > 0:
+ fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
+ bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -self.clamp_len, self.clamp_len)
+
+ if bsz is not None:
+ if bsz % 2 != 0:
+ raise ValueError(f"With bi_data, the batch size {bsz} should be divisible by 2")
+ fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2)
+ bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2)
+ else:
+ fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
+ bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)
+
+ pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
+ else:
+ fwd_pos_seq = tf.range(beg, end, -1.0)
+ if self.clamp_len > 0:
+ fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
+ pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz)
+
+ return pos_emb
+
+ @unpack_inputs
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ training: bool = False,
+ ):
+ if training and use_mems is None:
+ use_mems = self.use_mems_train
+ else:
+ use_mems = self.use_mems_eval
+
+ # the original code for XLNet uses shapes [len, bsz] with the batch dimension at the end
+ # but we want a unified interface in the library with the batch size on the first dimension
+ # so we move here the first dimension (batch) to the end
+
+ 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)[:2]
+ 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")
+
+ token_type_ids = tf.transpose(token_type_ids, perm=(1, 0)) if token_type_ids is not None else None
+ input_mask = tf.transpose(input_mask, perm=(1, 0)) if input_mask is not None else None
+ attention_mask = tf.transpose(attention_mask, perm=(1, 0)) if attention_mask is not None else None
+ perm_mask = tf.transpose(perm_mask, perm=(1, 2, 0)) if perm_mask is not None else None
+ target_mapping = tf.transpose(target_mapping, perm=(1, 2, 0)) if target_mapping is not None else None
+
+ mlen = shape_list(mems[0])[0] if mems is not None and mems[0] is not None else 0
+ klen = mlen + qlen
+
+ # Attention mask
+ # causal attention mask
+ if self.attn_type == "uni":
+ attn_mask = self.create_mask(qlen, mlen)
+ attn_mask = attn_mask[:, :, None, None]
+ elif self.attn_type == "bi":
+ attn_mask = None
+ else:
+ raise ValueError(f"Unsupported attention type: {self.attn_type}")
+
+ # data mask: input mask & perm mask
+ assert input_mask is None or attention_mask is None, (
+ "You can only use one of input_mask (uses 1 for padding) "
+ "or attention_mask (uses 0 for padding, added for compatibility with BERT). Please choose one."
+ )
+ if input_mask is None and attention_mask is not None:
+ one_cst = tf.constant(1.0)
+ input_mask = 1.0 - tf.cast(attention_mask, dtype=one_cst.dtype)
+ if input_mask is not None and perm_mask is not None:
+ data_mask = input_mask[None] + perm_mask
+ elif input_mask is not None and perm_mask is None:
+ data_mask = input_mask[None]
+ elif input_mask is None and perm_mask is not None:
+ data_mask = perm_mask
+ else:
+ data_mask = None
+
+ if data_mask is not None:
+ # all mems can be attended to
+ if mlen > 0:
+ mems_mask = tf.zeros([shape_list(data_mask)[0], mlen, bsz])
+ data_mask = tf.concat([mems_mask, data_mask], axis=1)
+ if attn_mask is None:
+ attn_mask = data_mask[:, :, :, None]
+ else:
+ attn_mask += data_mask[:, :, :, None]
+
+ if attn_mask is not None:
+ attn_mask = tf.cast(attn_mask > 0, dtype=attn_mask.dtype)
+
+ if attn_mask is not None:
+ non_tgt_mask = -tf.eye(qlen)
+ if mlen > 0:
+ non_tgt_mask = tf.concat([tf.zeros([qlen, mlen]), non_tgt_mask], axis=-1)
+ non_tgt_mask = tf.cast((attn_mask + non_tgt_mask[:, :, None, None]) > 0, dtype=non_tgt_mask.dtype)
+ else:
+ non_tgt_mask = None
+
+ # Word embeddings and prepare h & g hidden states
+ if inputs_embeds is not None:
+ word_emb_k = inputs_embeds
+ else:
+ check_embeddings_within_bounds(input_ids, self.word_embedding.vocab_size)
+ word_emb_k = self.word_embedding(input_ids)
+ output_h = self.dropout(word_emb_k, training=training)
+ if target_mapping is not None:
+ word_emb_q = tf.tile(self.mask_emb, [shape_list(target_mapping)[0], bsz, 1])
+ # else: # We removed the inp_q input which was same as target mapping
+ # inp_q_ext = inp_q[:, :, None]
+ # word_emb_q = inp_q_ext * self.mask_emb + (1 - inp_q_ext) * word_emb_k
+ output_g = self.dropout(word_emb_q, training=training)
+ else:
+ output_g = None
+
+ # Segment embedding
+ if token_type_ids is not None:
+ # Convert `token_type_ids` to one-hot `seg_mat`
+ if mlen > 0:
+ mem_pad = tf.zeros([mlen, bsz], dtype=token_type_ids.dtype)
+ cat_ids = tf.concat([mem_pad, token_type_ids], 0)
+ else:
+ cat_ids = token_type_ids
+
+ # `1` indicates not in the same segment [qlen x klen x bsz]
+ seg_mat = tf.cast(
+ tf.logical_not(tf.equal(token_type_ids[:, None], cat_ids[None, :])),
+ dtype=token_type_ids.dtype,
+ )
+ seg_mat = tf.one_hot(seg_mat, 2)
+ else:
+ seg_mat = None
+
+ # Positional encoding
+ pos_emb = self.relative_positional_encoding(qlen, klen, bsz=bsz)
+ pos_emb = self.dropout(pos_emb, training=training)
+
+ # 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
+
+ new_mems = ()
+ if mems is None:
+ mems = [None] * len(self.layer)
+
+ attentions = [] if output_attentions else None
+ hidden_states = [] if output_hidden_states else None
+ for i, layer_module in enumerate(self.layer):
+ # cache new mems
+ if use_mems:
+ new_mems = new_mems + (self.cache_mem(output_h, mems[i]),)
+ if output_hidden_states:
+ hidden_states.append((output_h, output_g) if output_g is not None else output_h)
+
+ outputs = layer_module(
+ output_h,
+ output_g,
+ non_tgt_mask,
+ attn_mask,
+ pos_emb,
+ seg_mat,
+ mems[i],
+ target_mapping,
+ head_mask[i],
+ output_attentions,
+ training=training,
+ )
+ output_h, output_g = outputs[:2]
+ if output_attentions:
+ attentions.append(outputs[2])
+
+ # Add last hidden state
+ if output_hidden_states:
+ hidden_states.append((output_h, output_g) if output_g is not None else output_h)
+
+ output = self.dropout(output_g if output_g is not None else output_h, training=training)
+
+ # Prepare outputs, we transpose back here to shape [bsz, len, hidden_dim] (cf. beginning of forward() method)
+ output = tf.transpose(output, perm=(1, 0, 2))
+
+ if not use_mems:
+ new_mems = None
+ if output_hidden_states:
+ if output_g is not None:
+ hidden_states = tuple(tf.transpose(h, perm=(1, 0, 2)) for hs in hidden_states for h in hs)
+ else:
+ hidden_states = tuple(tf.transpose(hs, perm=(1, 0, 2)) for hs in hidden_states)
+ if output_attentions:
+ if target_mapping is not None:
+ # when target_mapping is provided, there are 2-tuple of attentions
+ attentions = tuple(
+ tuple(tf.transpose(attn_stream, perm=(2, 3, 0, 1)) for attn_stream in t) for t in attentions
+ )
+ else:
+ attentions = tuple(tf.transpose(t, perm=(2, 3, 0, 1)) for t in attentions)
+
+ if not return_dict:
+ return tuple(v for v in [output, new_mems, hidden_states, attentions] if v is not None)
+
+ return TFXLNetModelOutput(
+ last_hidden_state=output, mems=new_mems, hidden_states=hidden_states, attentions=attentions
+ )
+
+
+class TFXLNetPreTrainedModel(TFPreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = XLNetConfig
+ base_model_prefix = "transformer"
+
+
+@dataclass
+class TFXLNetModelOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetModel`].
+
+ Args:
+ last_hidden_state (`tf.Tensor` of shape `(batch_size, num_predict, hidden_size)`):
+ Sequence of hidden-states at the last layer of the model.
+
+ `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict`
+ corresponds to `sequence_length`.
+ mems (`List[tf.Tensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFXLNetLMHeadModelOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetLMHeadModel`].
+
+ Args:
+ loss (`tf.Tensor` of shape *(1,)*, *optional*, returned when `labels` is provided)
+ Language modeling loss (for next-token prediction).
+ logits (`tf.Tensor` of shape `(batch_size, num_predict, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+
+ `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict`
+ corresponds to `sequence_length`.
+ mems (`List[tf.Tensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFXLNetForSequenceClassificationOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetForSequenceClassification`].
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `label` 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. 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: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFXLNetForTokenClassificationOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetForTokenClassificationOutput`].
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :
+ Classification loss.
+ logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.num_labels)`):
+ Classification scores (before SoftMax).
+ mems (`List[tf.Tensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFXLNetForMultipleChoiceOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetForMultipleChoice`].
+
+ Args:
+ loss (`tf.Tensor` of shape *(1,)*, *optional*, returned when `labels` is provided):
+ Classification loss.
+ logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
+ *num_choices* is the second dimension of the input tensors. (see *input_ids* above).
+
+ Classification scores (before SoftMax).
+ mems (`List[tf.Tensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+@dataclass
+class TFXLNetForQuestionAnsweringSimpleOutput(ModelOutput):
+ """
+ Output type of [`TFXLNetForQuestionAnsweringSimple`].
+
+ Args:
+ loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+ start_logits (`tf.Tensor` of shape `(batch_size, sequence_length,)`):
+ Span-start scores (before SoftMax).
+ end_logits (`tf.Tensor` of shape `(batch_size, sequence_length,)`):
+ Span-end scores (before SoftMax).
+ mems (`List[tf.Tensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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
+ start_logits: Optional[tf.Tensor] = None
+ end_logits: Optional[tf.Tensor] = None
+ mems: List[tf.Tensor] | None = None
+ hidden_states: Tuple[tf.Tensor, ...] | None = None
+ attentions: Tuple[tf.Tensor, ...] | None = None
+
+
+XLNET_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 ([`XLNetConfig`]): 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.
+"""
+
+XLNET_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)
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states (see `mems` 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.
+
+ `use_mems` has to be set to `True` to make use of `mems`.
+ perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*):
+ Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`:
+
+ - if `perm_mask[k, i, j] = 0`, i attend to j in batch k;
+ - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k.
+
+ If not set, each token attends to all the others (full bidirectional attention). Only used during
+ pretraining (to define factorization order) or for sequential decoding (generation).
+ target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*):
+ Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is
+ on the j-th token. Only used during pretraining for partial prediction or for sequential decoding
+ (generation).
+ 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)
+ input_mask (`torch.FloatTensor` of shape `{0}`, *optional*):
+ Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for
+ real tokens and 1 for padding which is kept for compatibility with the original code base.
+
+ Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **masked**,
+ - 0 for tokens that are **not masked**.
+
+ You can only uses one of `input_mask` and `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 XLNet Model transformer outputting raw hidden-states without any specific head on top.",
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetModel(TFXLNetPreTrainedModel):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFXLNetModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ training: bool = False,
+ ) -> Union[TFXLNetModelOutput, Tuple[tf.Tensor]]:
+ outputs = self.transformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ 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(
+ """
+ XLNet Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
+ """,
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetLMHeadModel(TFXLNetPreTrainedModel, TFCausalLanguageModelingLoss):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+ self.lm_loss = TFXLNetLMHead(config, self.transformer.word_embedding, name="lm_loss")
+ # generate fails to convert to a graph with XLNet
+ self.supports_xla_generation = False
+
+ def get_lm_head(self):
+ return self.lm_loss
+
+ 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.lm_loss.name
+
+ def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_mems=None, **kwargs):
+ # Add dummy token at the end (no attention on this one)
+ effective_batch_size = inputs.shape[0]
+ dummy_token = tf.zeros((effective_batch_size, 1), dtype=inputs.dtype)
+
+ # At every pass, the attention values for the new token and the two last generated tokens
+ # are computed, the rest is reloaded from the `past` cache. A purely auto-regressive model would have
+ # offset = 1; offset = 2 seems to have slightly better computation.
+ offset = 2
+
+ if past_key_values:
+ input_ids = tf.concat([inputs[:, -offset:], dummy_token], axis=1)
+ else:
+ input_ids = tf.concat([inputs, dummy_token], axis=1)
+
+ # Build permutation mask so that previous tokens don't see last token
+ sequence_length = input_ids.shape[1]
+ perm_mask = tf.zeros((effective_batch_size, sequence_length, sequence_length - 1))
+ perm_mask_seq_end = tf.ones((effective_batch_size, sequence_length, 1))
+ perm_mask = tf.concat([perm_mask, perm_mask_seq_end], axis=-1)
+
+ # We'll only predict the last token
+ target_mapping = tf.zeros((effective_batch_size, 1, sequence_length - 1))
+ target_mapping_seq_end = tf.ones((effective_batch_size, 1, 1))
+ target_mapping = tf.concat([target_mapping, target_mapping_seq_end], axis=-1)
+
+ inputs = {
+ "input_ids": input_ids,
+ "perm_mask": perm_mask,
+ "target_mapping": target_mapping,
+ "use_mems": use_mems,
+ }
+
+ # if past is defined in model kwargs then use it for faster decoding
+ if past_key_values:
+ inputs["mems"] = tuple(layer_past[:-offset, :, :] for layer_past in past_key_values)
+
+ return inputs
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=TFXLNetLMHeadModelOutput, config_class=_CONFIG_FOR_DOC)
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: 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: bool = False,
+ ) -> Union[TFXLNetLMHeadModelOutput, 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]`.
+
+ Return:
+
+ Examples:
+
+ ```python
+ >>> import tensorflow as tf
+ >>> import numpy as np
+ >>> from transformers import AutoTokenizer, TFXLNetLMHeadModel
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("xlnet/xlnet-large-cased")
+ >>> model = TFXLNetLMHeadModel.from_pretrained("xlnet/xlnet-large-cased")
+
+ >>> # We show how to setup inputs to predict a next token using a bi-directional context.
+ >>> input_ids = tf.constant(tokenizer.encode("Hello, my dog is very ", add_special_tokens=True))[
+ ... None, :
+ ... ] # We will predict the masked token
+
+ >>> perm_mask = np.zeros((1, input_ids.shape[1], input_ids.shape[1]))
+ >>> perm_mask[:, :, -1] = 1.0 # Previous tokens don't see last token
+
+ >>> target_mapping = np.zeros(
+ ... (1, 1, input_ids.shape[1])
+ ... ) # Shape [1, 1, seq_length] => let's predict one token
+ >>> target_mapping[
+ ... 0, 0, -1
+ ... ] = 1.0 # Our first (and only) prediction will be the last token of the sequence (the masked token)
+
+ >>> outputs = model(
+ ... input_ids,
+ ... perm_mask=tf.constant(perm_mask, dtype=tf.float32),
+ ... target_mapping=tf.constant(target_mapping, dtype=tf.float32),
+ ... )
+
+ >>> next_token_logits = outputs[
+ ... 0
+ ... ] # Output has shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size]
+ ```"""
+ transformer_outputs = self.transformer(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ training=training,
+ )
+ hidden_state = transformer_outputs[0]
+ logits = self.lm_loss(hidden_state, training=training)
+
+ loss = None
+ if labels is not None:
+ loss = 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 TFXLNetLMHeadModelOutput(
+ loss=loss,
+ logits=logits,
+ mems=transformer_outputs.mems,
+ 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_loss", None) is not None:
+ with tf.name_scope(self.lm_loss.name):
+ self.lm_loss.build(None)
+
+
+@add_start_docstrings(
+ """
+ XLNet Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g.
+ for GLUE tasks.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetForSequenceClassification(TFXLNetPreTrainedModel, TFSequenceClassificationLoss):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.num_labels = config.num_labels
+
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+ self.sequence_summary = TFSequenceSummary(
+ config, initializer_range=config.initializer_range, name="sequence_summary"
+ )
+ self.logits_proj = keras.layers.Dense(
+ config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="logits_proj"
+ )
+ self.config = config
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFXLNetForSequenceClassificationOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: 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: bool = False,
+ ) -> Union[TFXLNetForSequenceClassificationOutput, 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ training=training,
+ )
+ output = transformer_outputs[0]
+
+ output = self.sequence_summary(output)
+ logits = self.logits_proj(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 TFXLNetForSequenceClassificationOutput(
+ loss=loss,
+ logits=logits,
+ mems=transformer_outputs.mems,
+ 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.d_model])
+
+
+@add_start_docstrings(
+ """
+ XLNET 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.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetForMultipleChoice(TFXLNetPreTrainedModel, TFMultipleChoiceLoss):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+ self.sequence_summary = TFSequenceSummary(
+ config, initializer_range=config.initializer_range, name="sequence_summary"
+ )
+ self.logits_proj = keras.layers.Dense(
+ 1, kernel_initializer=get_initializer(config.initializer_range), name="logits_proj"
+ )
+ self.config = config
+
+ @unpack_inputs
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFXLNetForMultipleChoiceOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: 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: bool = False,
+ ) -> Union[TFXLNetForMultipleChoiceOutput, 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_input_mask = tf.reshape(input_mask, (-1, seq_length)) if input_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
+ )
+ transformer_outputs = self.transformer(
+ flat_input_ids,
+ flat_attention_mask,
+ mems,
+ perm_mask,
+ target_mapping,
+ flat_token_type_ids,
+ flat_input_mask,
+ head_mask,
+ flat_inputs_embeds,
+ use_mems,
+ 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 TFXLNetForMultipleChoiceOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ mems=transformer_outputs.mems,
+ 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.d_model])
+
+
+@add_start_docstrings(
+ """
+ XLNet 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.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetForTokenClassification(TFXLNetPreTrainedModel, TFTokenClassificationLoss):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.num_labels = config.num_labels
+
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+ 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(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFXLNetForTokenClassificationOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: 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: bool = False,
+ ) -> Union[TFXLNetForTokenClassificationOutput, 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ training=training,
+ )
+ output = transformer_outputs[0]
+ logits = self.classifier(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 TFXLNetForTokenClassificationOutput(
+ loss=loss,
+ logits=logits,
+ mems=transformer_outputs.mems,
+ 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(
+ """
+ XLNet 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`).
+ """,
+ XLNET_START_DOCSTRING,
+)
+class TFXLNetForQuestionAnsweringSimple(TFXLNetPreTrainedModel, TFQuestionAnsweringLoss):
+ def __init__(self, config, *inputs, **kwargs):
+ super().__init__(config, *inputs, **kwargs)
+ self.transformer = TFXLNetMainLayer(config, name="transformer")
+ 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(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=TFXLNetForQuestionAnsweringSimpleOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def call(
+ self,
+ input_ids: TFModelInputType | None = None,
+ attention_mask: np.ndarray | tf.Tensor | None = None,
+ mems: np.ndarray | tf.Tensor | None = None,
+ perm_mask: np.ndarray | tf.Tensor | None = None,
+ target_mapping: np.ndarray | tf.Tensor | None = None,
+ token_type_ids: np.ndarray | tf.Tensor | None = None,
+ input_mask: np.ndarray | tf.Tensor | None = None,
+ head_mask: np.ndarray | tf.Tensor | None = None,
+ inputs_embeds: np.ndarray | tf.Tensor | None = None,
+ use_mems: Optional[bool] = 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[TFXLNetForQuestionAnsweringSimpleOutput, 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ 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 TFXLNetForQuestionAnsweringSimpleOutput(
+ loss=loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ mems=transformer_outputs.mems,
+ 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])
+
+
+__all__ = [
+ "TFXLNetForMultipleChoice",
+ "TFXLNetForQuestionAnsweringSimple",
+ "TFXLNetForSequenceClassification",
+ "TFXLNetForTokenClassification",
+ "TFXLNetLMHeadModel",
+ "TFXLNetMainLayer",
+ "TFXLNetModel",
+ "TFXLNetPreTrainedModel",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/modeling_xlnet.py b/lib/python3.12/site-packages/transformers/models/xlnet/modeling_xlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..db0703f949c3870b3f49c861c542c4bbc64d36d1
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/modeling_xlnet.py
@@ -0,0 +1,2097 @@
+# 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 XLNet model.
+"""
+
+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 ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_utils import PoolerAnswerClass, PoolerEndLogits, PoolerStartLogits, PreTrainedModel, SequenceSummary
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from .configuration_xlnet import XLNetConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "xlnet/xlnet-base-cased"
+_CONFIG_FOR_DOC = "XLNetConfig"
+
+
+def build_tf_xlnet_to_pytorch_map(model, config, tf_weights=None):
+ """
+ A map of modules from TF to PyTorch. 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"):
+ if hasattr(model, "lm_loss"):
+ # We will load also the output bias
+ tf_to_pt_map["model/lm_loss/bias"] = model.lm_loss.bias
+ if hasattr(model, "sequence_summary") and "model/sequnece_summary/summary/kernel" in tf_weights:
+ # We will load also the sequence summary
+ tf_to_pt_map["model/sequnece_summary/summary/kernel"] = model.sequence_summary.summary.weight
+ tf_to_pt_map["model/sequnece_summary/summary/bias"] = model.sequence_summary.summary.bias
+ if (
+ hasattr(model, "logits_proj")
+ and config.finetuning_task is not None
+ and f"model/regression_{config.finetuning_task}/logit/kernel" in tf_weights
+ ):
+ tf_to_pt_map[f"model/regression_{config.finetuning_task}/logit/kernel"] = model.logits_proj.weight
+ tf_to_pt_map[f"model/regression_{config.finetuning_task}/logit/bias"] = model.logits_proj.bias
+
+ # Now load the rest of the transformer
+ model = model.transformer
+
+ # Embeddings and output
+ tf_to_pt_map.update(
+ {
+ "model/transformer/word_embedding/lookup_table": model.word_embedding.weight,
+ "model/transformer/mask_emb/mask_emb": model.mask_emb,
+ }
+ )
+
+ # Transformer blocks
+ for i, b in enumerate(model.layer):
+ layer_str = f"model/transformer/layer_{i}/"
+ tf_to_pt_map.update(
+ {
+ layer_str + "rel_attn/LayerNorm/gamma": b.rel_attn.layer_norm.weight,
+ layer_str + "rel_attn/LayerNorm/beta": b.rel_attn.layer_norm.bias,
+ layer_str + "rel_attn/o/kernel": b.rel_attn.o,
+ layer_str + "rel_attn/q/kernel": b.rel_attn.q,
+ layer_str + "rel_attn/k/kernel": b.rel_attn.k,
+ layer_str + "rel_attn/r/kernel": b.rel_attn.r,
+ layer_str + "rel_attn/v/kernel": b.rel_attn.v,
+ layer_str + "ff/LayerNorm/gamma": b.ff.layer_norm.weight,
+ layer_str + "ff/LayerNorm/beta": b.ff.layer_norm.bias,
+ layer_str + "ff/layer_1/kernel": b.ff.layer_1.weight,
+ layer_str + "ff/layer_1/bias": b.ff.layer_1.bias,
+ layer_str + "ff/layer_2/kernel": b.ff.layer_2.weight,
+ layer_str + "ff/layer_2/bias": b.ff.layer_2.bias,
+ }
+ )
+
+ # Relative positioning biases
+ if config.untie_r:
+ r_r_list = []
+ r_w_list = []
+ r_s_list = []
+ seg_embed_list = []
+ for b in model.layer:
+ r_r_list.append(b.rel_attn.r_r_bias)
+ r_w_list.append(b.rel_attn.r_w_bias)
+ r_s_list.append(b.rel_attn.r_s_bias)
+ seg_embed_list.append(b.rel_attn.seg_embed)
+ else:
+ r_r_list = [model.r_r_bias]
+ r_w_list = [model.r_w_bias]
+ r_s_list = [model.r_s_bias]
+ seg_embed_list = [model.seg_embed]
+ tf_to_pt_map.update(
+ {
+ "model/transformer/r_r_bias": r_r_list,
+ "model/transformer/r_w_bias": r_w_list,
+ "model/transformer/r_s_bias": r_s_list,
+ "model/transformer/seg_embed": seg_embed_list,
+ }
+ )
+ return tf_to_pt_map
+
+
+def load_tf_weights_in_xlnet(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
+ # 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
+
+ # Build TF to PyTorch weights loading map
+ tf_to_pt_map = build_tf_xlnet_to_pytorch_map(model, config, tf_weights)
+
+ for name, pointer in tf_to_pt_map.items():
+ logger.info(f"Importing {name}")
+ if name not in tf_weights:
+ logger.info(f"{name} not in tf pre-trained weights, skipping")
+ continue
+ 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 and ("ff" in name or "summary" in name or "logit" in name):
+ logger.info("Transposing")
+ array = np.transpose(array)
+ if isinstance(pointer, list):
+ # Here we will split the TF weights
+ assert len(pointer) == array.shape[0], (
+ f"Pointer length {len(pointer)} and array length {array.shape[0]} mismatched"
+ )
+ for i, p_i in enumerate(pointer):
+ arr_i = array[i, ...]
+ try:
+ assert p_i.shape == arr_i.shape, (
+ f"Pointer shape {p_i.shape} and array shape {arr_i.shape} mismatched"
+ )
+ 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 XLNetRelativeAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ if config.d_model % config.n_head != 0:
+ raise ValueError(
+ f"The hidden size ({config.d_model}) is not a multiple of the number of attention "
+ f"heads ({config.n_head}"
+ )
+
+ self.n_head = config.n_head
+ self.d_head = config.d_head
+ self.d_model = config.d_model
+ self.scale = 1 / (config.d_head**0.5)
+
+ self.q = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head))
+ self.k = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head))
+ self.v = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head))
+ self.o = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head))
+ self.r = nn.Parameter(torch.FloatTensor(config.d_model, self.n_head, self.d_head))
+
+ self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
+ self.r_s_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))
+ self.seg_embed = nn.Parameter(torch.FloatTensor(2, self.n_head, self.d_head))
+
+ self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.dropout)
+
+ def prune_heads(self, heads):
+ raise NotImplementedError
+
+ @staticmethod
+ def rel_shift(x, klen=-1):
+ """perform relative shift to form the relative attention score."""
+ x_size = x.shape
+
+ x = x.reshape(x_size[1], x_size[0], x_size[2], x_size[3])
+ x = x[1:, ...]
+ x = x.reshape(x_size[0], x_size[1] - 1, x_size[2], x_size[3])
+ # x = x[:, 0:klen, :, :]
+ x = torch.index_select(x, 1, torch.arange(klen, device=x.device, dtype=torch.long))
+
+ return x
+
+ @staticmethod
+ def rel_shift_bnij(x, klen=-1):
+ x_size = x.shape
+
+ x = x.reshape(x_size[0], x_size[1], x_size[3], x_size[2])
+ x = x[:, :, 1:, :]
+ x = x.reshape(x_size[0], x_size[1], x_size[2], x_size[3] - 1)
+ # Note: the tensor-slice form was faster in my testing than torch.index_select
+ # However, tracing doesn't like the nature of the slice, and if klen changes
+ # during the run then it'll fail, whereas index_select will be fine.
+ x = torch.index_select(x, 3, torch.arange(klen, device=x.device, dtype=torch.long))
+ # x = x[:, :, :, :klen]
+
+ return x
+
+ def rel_attn_core(
+ self,
+ q_head,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat=None,
+ attn_mask=None,
+ head_mask=None,
+ output_attentions=False,
+ ):
+ """Core relative positional attention operations."""
+
+ # content based attention score
+ ac = torch.einsum("ibnd,jbnd->bnij", q_head + self.r_w_bias, k_head_h)
+
+ # position based attention score
+ bd = torch.einsum("ibnd,jbnd->bnij", q_head + self.r_r_bias, k_head_r)
+ bd = self.rel_shift_bnij(bd, klen=ac.shape[3])
+
+ # segment based attention score
+ if seg_mat is None:
+ ef = 0
+ else:
+ ef = torch.einsum("ibnd,snd->ibns", q_head + self.r_s_bias, self.seg_embed)
+ ef = torch.einsum("ijbs,ibns->bnij", seg_mat, ef)
+
+ # merge attention scores and perform masking
+ attn_score = (ac + bd + ef) * self.scale
+ if attn_mask is not None:
+ # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
+ if attn_mask.dtype == torch.float16:
+ attn_score = attn_score - 65500 * torch.einsum("ijbn->bnij", attn_mask)
+ else:
+ attn_score = attn_score - 1e30 * torch.einsum("ijbn->bnij", attn_mask)
+
+ # attention probability
+ attn_prob = nn.functional.softmax(attn_score, dim=3)
+ attn_prob = self.dropout(attn_prob)
+
+ # Mask heads if we want to
+ if head_mask is not None:
+ attn_prob = attn_prob * torch.einsum("ijbn->bnij", head_mask)
+
+ # attention output
+ attn_vec = torch.einsum("bnij,jbnd->ibnd", attn_prob, v_head_h)
+
+ if output_attentions:
+ return attn_vec, torch.einsum("bnij->ijbn", attn_prob)
+
+ return attn_vec
+
+ def post_attention(self, h, attn_vec, residual=True):
+ """Post-attention processing."""
+ # post-attention projection (back to `d_model`)
+ attn_out = torch.einsum("ibnd,hnd->ibh", attn_vec, self.o)
+
+ attn_out = self.dropout(attn_out)
+ if residual:
+ attn_out = attn_out + h
+ output = self.layer_norm(attn_out)
+
+ return output
+
+ def forward(
+ self,
+ h,
+ g,
+ attn_mask_h,
+ attn_mask_g,
+ r,
+ seg_mat,
+ mems=None,
+ target_mapping=None,
+ head_mask=None,
+ output_attentions=False,
+ ):
+ if g is not None:
+ # Two-stream attention with relative positional encoding.
+ # content based attention score
+ if mems is not None and mems.dim() > 1:
+ cat = torch.cat([mems, h], dim=0)
+ else:
+ cat = h
+
+ # content-based key head
+ k_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.k)
+
+ # content-based value head
+ v_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.v)
+
+ # position-based key head
+ k_head_r = torch.einsum("ibh,hnd->ibnd", r, self.r)
+
+ # h-stream
+ # content-stream query head
+ q_head_h = torch.einsum("ibh,hnd->ibnd", h, self.q)
+
+ # core attention ops
+ attn_vec_h = self.rel_attn_core(
+ q_head_h,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat=seg_mat,
+ attn_mask=attn_mask_h,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ )
+
+ if output_attentions:
+ attn_vec_h, attn_prob_h = attn_vec_h
+
+ # post processing
+ output_h = self.post_attention(h, attn_vec_h)
+
+ # g-stream
+ # query-stream query head
+ q_head_g = torch.einsum("ibh,hnd->ibnd", g, self.q)
+
+ # core attention ops
+ if target_mapping is not None:
+ q_head_g = torch.einsum("mbnd,mlb->lbnd", q_head_g, target_mapping)
+ attn_vec_g = self.rel_attn_core(
+ q_head_g,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat=seg_mat,
+ attn_mask=attn_mask_g,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ )
+
+ if output_attentions:
+ attn_vec_g, attn_prob_g = attn_vec_g
+
+ attn_vec_g = torch.einsum("lbnd,mlb->mbnd", attn_vec_g, target_mapping)
+ else:
+ attn_vec_g = self.rel_attn_core(
+ q_head_g,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat=seg_mat,
+ attn_mask=attn_mask_g,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ )
+
+ if output_attentions:
+ attn_vec_g, attn_prob_g = attn_vec_g
+
+ # post processing
+ output_g = self.post_attention(g, attn_vec_g)
+
+ if output_attentions:
+ attn_prob = attn_prob_h, attn_prob_g
+
+ else:
+ # Multi-head attention with relative positional encoding
+ if mems is not None and mems.dim() > 1:
+ cat = torch.cat([mems, h], dim=0)
+ else:
+ cat = h
+
+ # content heads
+ q_head_h = torch.einsum("ibh,hnd->ibnd", h, self.q)
+ k_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.k)
+ v_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.v)
+
+ # positional heads
+ # type casting for fp16 support
+ k_head_r = torch.einsum("ibh,hnd->ibnd", r.type(self.r.dtype), self.r)
+
+ # core attention ops
+ attn_vec = self.rel_attn_core(
+ q_head_h,
+ k_head_h,
+ v_head_h,
+ k_head_r,
+ seg_mat=seg_mat,
+ attn_mask=attn_mask_h,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ )
+
+ if output_attentions:
+ attn_vec, attn_prob = attn_vec
+
+ # post processing
+ output_h = self.post_attention(h, attn_vec)
+ output_g = None
+
+ outputs = (output_h, output_g)
+ if output_attentions:
+ outputs = outputs + (attn_prob,)
+ return outputs
+
+
+class XLNetFeedForward(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+ self.layer_1 = nn.Linear(config.d_model, config.d_inner)
+ self.layer_2 = nn.Linear(config.d_inner, config.d_model)
+ self.dropout = nn.Dropout(config.dropout)
+ if isinstance(config.ff_activation, str):
+ self.activation_function = ACT2FN[config.ff_activation]
+ else:
+ self.activation_function = config.ff_activation
+
+ def forward(self, inp):
+ output = inp
+ output = self.layer_1(output)
+ output = self.activation_function(output)
+ output = self.dropout(output)
+ output = self.layer_2(output)
+ output = self.dropout(output)
+ output = self.layer_norm(output + inp)
+ return output
+
+
+class XLNetLayer(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.rel_attn = XLNetRelativeAttention(config)
+ self.ff = XLNetFeedForward(config)
+ self.dropout = nn.Dropout(config.dropout)
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+
+ def forward(
+ self,
+ output_h,
+ output_g,
+ attn_mask_h,
+ attn_mask_g,
+ r,
+ seg_mat,
+ mems=None,
+ target_mapping=None,
+ head_mask=None,
+ output_attentions=False,
+ ):
+ outputs = self.rel_attn(
+ output_h,
+ output_g,
+ attn_mask_h,
+ attn_mask_g,
+ r,
+ seg_mat,
+ mems=mems,
+ target_mapping=target_mapping,
+ head_mask=head_mask,
+ output_attentions=output_attentions,
+ )
+ output_h, output_g = outputs[:2]
+
+ if output_g is not None:
+ output_g = apply_chunking_to_forward(
+ self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_g
+ )
+ output_h = apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_h)
+
+ outputs = (output_h, output_g) + outputs[2:] # Add again attentions if there are there
+ return outputs
+
+ def ff_chunk(self, output_x):
+ output_x = self.ff(output_x)
+ return output_x
+
+
+class XLNetPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = XLNetConfig
+ load_tf_weights = load_tf_weights_in_xlnet
+ base_model_prefix = "transformer"
+
+ 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)
+ elif isinstance(module, XLNetRelativeAttention):
+ for param in [
+ module.q,
+ module.k,
+ module.v,
+ module.o,
+ module.r,
+ module.r_r_bias,
+ module.r_s_bias,
+ module.r_w_bias,
+ module.seg_embed,
+ ]:
+ param.data.normal_(mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, XLNetModel):
+ module.mask_emb.data.normal_(mean=0.0, std=self.config.initializer_range)
+
+
+@dataclass
+class XLNetModelOutput(ModelOutput):
+ """
+ Output type of [`XLNetModel`].
+
+ Args:
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_predict, hidden_size)`):
+ Sequence of hidden-states at the last layer of the model.
+
+ `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict`
+ corresponds to `sequence_length`.
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetLMHeadModelOutput(ModelOutput):
+ """
+ Output type of [`XLNetLMHeadModel`].
+
+ Args:
+ loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided)
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, num_predict, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+
+ `num_predict` corresponds to `target_mapping.shape[1]`. If `target_mapping` is `None`, then `num_predict`
+ corresponds to `sequence_length`.
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[torch.FloatTensor] = None
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetForSequenceClassificationOutput(ModelOutput):
+ """
+ Output type of [`XLNetForSequenceClassification`].
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `label` 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. 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: Optional[torch.FloatTensor] = None
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetForTokenClassificationOutput(ModelOutput):
+ """
+ Output type of [`XLNetForTokenClassificationOutput`].
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :
+ Classification loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):
+ Classification scores (before SoftMax).
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[torch.FloatTensor] = None
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetForMultipleChoiceOutput(ModelOutput):
+ """
+ Output type of [`XLNetForMultipleChoice`].
+
+ Args:
+ loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided):
+ Classification loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`):
+ *num_choices* is the second dimension of the input tensors. (see *input_ids* above).
+
+ Classification scores (before SoftMax).
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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: Optional[torch.FloatTensor] = None
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetForQuestionAnsweringSimpleOutput(ModelOutput):
+ """
+ Output type of [`XLNetForQuestionAnsweringSimple`].
+
+ Args:
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+ start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length,)`):
+ Span-start scores (before SoftMax).
+ end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length,)`):
+ Span-end scores (before SoftMax).
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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
+ start_logits: Optional[torch.FloatTensor] = None
+ end_logits: Optional[torch.FloatTensor] = None
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class XLNetForQuestionAnsweringOutput(ModelOutput):
+ """
+ Output type of [`XLNetForQuestionAnswering`].
+
+ 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.
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states. 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
+ 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
+ mems: Optional[List[torch.FloatTensor]] = None
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+XLNET_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 ([`XLNetConfig`]): 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.
+"""
+
+XLNET_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)
+ mems (`List[torch.FloatTensor]` of length `config.n_layers`):
+ Contains pre-computed hidden-states (see `mems` 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.
+
+ `use_mems` has to be set to `True` to make use of `mems`.
+ perm_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length)`, *optional*):
+ Mask to indicate the attention pattern for each input token with values selected in `[0, 1]`:
+
+ - if `perm_mask[k, i, j] = 0`, i attend to j in batch k;
+ - if `perm_mask[k, i, j] = 1`, i does not attend to j in batch k.
+
+ If not set, each token attends to all the others (full bidirectional attention). Only used during
+ pretraining (to define factorization order) or for sequential decoding (generation).
+ target_mapping (`torch.FloatTensor` of shape `(batch_size, num_predict, sequence_length)`, *optional*):
+ Mask to indicate the output tokens to use. If `target_mapping[k, i, j] = 1`, the i-th predict in batch k is
+ on the j-th token. Only used during pretraining for partial prediction or for sequential decoding
+ (generation).
+ 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)
+ input_mask (`torch.FloatTensor` of shape `{0}`, *optional*):
+ Mask to avoid performing attention on padding token indices. Negative of `attention_mask`, i.e. with 0 for
+ real tokens and 1 for padding which is kept for compatibility with the original code base.
+
+ Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **masked**,
+ - 0 for tokens that are **not masked**.
+
+ You can only uses one of `input_mask` and `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 XLNet Model transformer outputting raw hidden-states without any specific head on top.",
+ XLNET_START_DOCSTRING,
+)
+class XLNetModel(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.mem_len = config.mem_len
+ self.reuse_len = config.reuse_len
+ self.d_model = config.d_model
+ self.same_length = config.same_length
+ self.attn_type = config.attn_type
+ self.bi_data = config.bi_data
+ self.clamp_len = config.clamp_len
+ self.n_layer = config.n_layer
+
+ self.word_embedding = nn.Embedding(config.vocab_size, config.d_model)
+ self.mask_emb = nn.Parameter(torch.FloatTensor(1, 1, config.d_model))
+ self.layer = nn.ModuleList([XLNetLayer(config) for _ in range(config.n_layer)])
+ self.dropout = nn.Dropout(config.dropout)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.word_embedding
+
+ def set_input_embeddings(self, new_embeddings):
+ self.word_embedding = new_embeddings
+
+ def _prune_heads(self, heads_to_prune):
+ raise NotImplementedError
+
+ def create_mask(self, qlen, mlen):
+ """
+ Creates causal attention mask. Float mask where 1.0 indicates masked, 0.0 indicates not-masked.
+
+ Args:
+ qlen: Sequence length
+ mlen: Mask length
+
+ ::
+
+ same_length=False: same_length=True: < qlen > < qlen >
+ ^ [0 0 0 0 0 1 1 1 1] [0 0 0 0 0 1 1 1 1]
+ [0 0 0 0 0 0 1 1 1] [1 0 0 0 0 0 1 1 1]
+ qlen [0 0 0 0 0 0 0 1 1] [1 1 0 0 0 0 0 1 1]
+ [0 0 0 0 0 0 0 0 1] [1 1 1 0 0 0 0 0 1]
+ v [0 0 0 0 0 0 0 0 0] [1 1 1 1 0 0 0 0 0]
+
+ """
+ mask = torch.ones((qlen, qlen + mlen), device=self.device)
+ if self.same_length:
+ mask_lo = mask[:, :qlen].tril(-1)
+ mask.triu_(mlen + 1)
+ mask[:, :qlen] += mask_lo
+ else:
+ mask.triu_(mlen + 1)
+
+ return mask
+
+ def cache_mem(self, curr_out, prev_mem):
+ # cache hidden states into memory.
+ if self.reuse_len is not None and self.reuse_len > 0:
+ curr_out = curr_out[: self.reuse_len]
+
+ if self.mem_len is None or self.mem_len == 0:
+ # If `use_mems` is active but no `mem_len` is defined, the model behaves like GPT-2 at inference time
+ # and returns all of the past and current hidden states.
+ cutoff = 0
+ else:
+ # If `use_mems` is active and `mem_len` is defined, the model returns the last `mem_len` hidden
+ # states. This is the preferred setting for training and long-form generation.
+ cutoff = -self.mem_len
+ if prev_mem is None:
+ # if `use_mems` is active and `mem_len` is defined, the model
+ new_mem = curr_out[cutoff:]
+ else:
+ new_mem = torch.cat([prev_mem, curr_out], dim=0)[cutoff:]
+
+ return new_mem.detach()
+
+ @staticmethod
+ def positional_embedding(pos_seq, inv_freq, bsz=None):
+ sinusoid_inp = torch.einsum("i,d->id", pos_seq, inv_freq)
+ pos_emb = torch.cat([torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)], dim=-1)
+ pos_emb = pos_emb[:, None, :]
+
+ if bsz is not None:
+ pos_emb = pos_emb.expand(-1, bsz, -1)
+
+ return pos_emb
+
+ def relative_positional_encoding(self, qlen, klen, bsz=None):
+ # create relative positional encoding.
+ freq_seq = torch.arange(0, self.d_model, 2.0, dtype=torch.int64).float()
+ inv_freq = 1 / torch.pow(10000, (freq_seq / self.d_model))
+
+ if self.attn_type == "bi":
+ # beg, end = klen - 1, -qlen
+ beg, end = klen, -qlen
+ elif self.attn_type == "uni":
+ # beg, end = klen - 1, -1
+ beg, end = klen, -1
+ else:
+ raise ValueError(f"Unknown `attn_type` {self.attn_type}.")
+
+ if self.bi_data:
+ fwd_pos_seq = torch.arange(beg, end, -1.0, dtype=torch.int64).float()
+ bwd_pos_seq = torch.arange(-beg, -end, 1.0, dtype=torch.int64).float()
+
+ if self.clamp_len > 0:
+ fwd_pos_seq = fwd_pos_seq.clamp(-self.clamp_len, self.clamp_len)
+ bwd_pos_seq = bwd_pos_seq.clamp(-self.clamp_len, self.clamp_len)
+
+ if bsz is not None:
+ fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2)
+ bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2)
+ else:
+ fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
+ bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)
+
+ pos_emb = torch.cat([fwd_pos_emb, bwd_pos_emb], dim=1)
+ else:
+ fwd_pos_seq = torch.arange(beg, end, -1.0, dtype=torch.int64).float()
+ if self.clamp_len > 0:
+ fwd_pos_seq = fwd_pos_seq.clamp(-self.clamp_len, self.clamp_len)
+ pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz)
+
+ return pos_emb
+
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=XLNetModelOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_mask: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete after depreciation warning is removed
+ ) -> Union[Tuple, XLNetModelOutput]:
+ 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 "use_cache" in kwargs:
+ warnings.warn(
+ "The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems`"
+ " instead.",
+ FutureWarning,
+ )
+ use_mems = kwargs["use_cache"]
+
+ if self.training:
+ use_mems = use_mems if use_mems is not None else self.config.use_mems_train
+ else:
+ use_mems = use_mems if use_mems is not None else self.config.use_mems_eval
+
+ # the original code for XLNet uses shapes [len, bsz] with the batch dimension at the end
+ # but we want a unified interface in the library with the batch size on the first dimension
+ # so we move here the first dimension (batch) to the end
+ 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.shape[0], input_ids.shape[1]
+ 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")
+
+ token_type_ids = token_type_ids.transpose(0, 1).contiguous() if token_type_ids is not None else None
+ input_mask = input_mask.transpose(0, 1).contiguous() if input_mask is not None else None
+ attention_mask = attention_mask.transpose(0, 1).contiguous() if attention_mask is not None else None
+ perm_mask = perm_mask.permute(1, 2, 0).contiguous() if perm_mask is not None else None
+ target_mapping = target_mapping.permute(1, 2, 0).contiguous() if target_mapping is not None else None
+
+ mlen = mems[0].shape[0] if mems is not None and mems[0] is not None else 0
+ klen = mlen + qlen
+
+ dtype_float = self.dtype
+ device = self.device
+
+ # Attention mask
+ # causal attention mask
+ if self.attn_type == "uni":
+ attn_mask = self.create_mask(qlen, mlen)
+ attn_mask = attn_mask[:, :, None, None]
+ elif self.attn_type == "bi":
+ attn_mask = None
+ else:
+ raise ValueError(f"Unsupported attention type: {self.attn_type}")
+
+ # data mask: input mask & perm mask
+ assert input_mask is None or attention_mask is None, "You can only use one of input_mask (uses 1 for padding) "
+ "or attention_mask (uses 0 for padding, added for compatibility with BERT). Please choose one."
+ if input_mask is None and attention_mask is not None:
+ input_mask = 1.0 - attention_mask
+ if input_mask is not None and perm_mask is not None:
+ data_mask = input_mask[None] + perm_mask
+ elif input_mask is not None and perm_mask is None:
+ data_mask = input_mask[None]
+ elif input_mask is None and perm_mask is not None:
+ data_mask = perm_mask
+ else:
+ data_mask = None
+
+ if data_mask is not None:
+ # all mems can be attended to
+ if mlen > 0:
+ mems_mask = torch.zeros([data_mask.shape[0], mlen, bsz]).to(data_mask)
+ data_mask = torch.cat([mems_mask, data_mask], dim=1)
+ if attn_mask is None:
+ attn_mask = data_mask[:, :, :, None]
+ else:
+ attn_mask += data_mask[:, :, :, None]
+
+ if attn_mask is not None:
+ attn_mask = (attn_mask > 0).to(dtype_float)
+
+ if attn_mask is not None:
+ non_tgt_mask = -torch.eye(qlen).to(attn_mask)
+ if mlen > 0:
+ non_tgt_mask = torch.cat([torch.zeros([qlen, mlen]).to(attn_mask), non_tgt_mask], dim=-1)
+ non_tgt_mask = ((attn_mask + non_tgt_mask[:, :, None, None]) > 0).to(attn_mask)
+ else:
+ non_tgt_mask = None
+
+ # Word embeddings and prepare h & g hidden states
+ if inputs_embeds is not None:
+ word_emb_k = inputs_embeds
+ else:
+ word_emb_k = self.word_embedding(input_ids)
+ output_h = self.dropout(word_emb_k)
+ if target_mapping is not None:
+ word_emb_q = self.mask_emb.expand(target_mapping.shape[0], bsz, -1)
+ # else: # We removed the inp_q input which was same as target mapping
+ # inp_q_ext = inp_q[:, :, None]
+ # word_emb_q = inp_q_ext * self.mask_emb + (1 - inp_q_ext) * word_emb_k
+ output_g = self.dropout(word_emb_q)
+ else:
+ output_g = None
+
+ # Segment embedding
+ if token_type_ids is not None:
+ # Convert `token_type_ids` to one-hot `seg_mat`
+ if mlen > 0:
+ mem_pad = torch.zeros([mlen, bsz], dtype=torch.long, device=device)
+ cat_ids = torch.cat([mem_pad, token_type_ids], dim=0)
+ else:
+ cat_ids = token_type_ids
+
+ # `1` indicates not in the same segment [qlen x klen x bsz]
+ seg_mat = (token_type_ids[:, None] != cat_ids[None, :]).long()
+ seg_mat = nn.functional.one_hot(seg_mat, num_classes=2).to(dtype_float)
+ else:
+ seg_mat = None
+
+ # Positional encoding
+ pos_emb = self.relative_positional_encoding(qlen, klen, bsz=bsz)
+ pos_emb = pos_emb.to(output_h.device)
+ pos_emb = self.dropout(pos_emb)
+
+ # 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
+
+ new_mems = ()
+ if mems is None:
+ mems = [None] * len(self.layer)
+
+ attentions = [] if output_attentions else None
+ hidden_states = [] if output_hidden_states else None
+ for i, layer_module in enumerate(self.layer):
+ if use_mems:
+ # cache new mems
+ new_mems = new_mems + (self.cache_mem(output_h, mems[i]),)
+ if output_hidden_states:
+ hidden_states.append((output_h, output_g) if output_g is not None else output_h)
+
+ outputs = layer_module(
+ output_h,
+ output_g,
+ attn_mask_h=non_tgt_mask,
+ attn_mask_g=attn_mask,
+ r=pos_emb,
+ seg_mat=seg_mat,
+ mems=mems[i],
+ target_mapping=target_mapping,
+ head_mask=head_mask[i],
+ output_attentions=output_attentions,
+ )
+ output_h, output_g = outputs[:2]
+ if output_attentions:
+ attentions.append(outputs[2])
+
+ # Add last hidden state
+ if output_hidden_states:
+ hidden_states.append((output_h, output_g) if output_g is not None else output_h)
+
+ output = self.dropout(output_g if output_g is not None else output_h)
+
+ # Prepare outputs, we transpose back here to shape [bsz, len, hidden_dim] (cf. beginning of forward() method)
+ output = output.permute(1, 0, 2).contiguous()
+
+ if not use_mems:
+ new_mems = None
+
+ if output_hidden_states:
+ if output_g is not None:
+ hidden_states = tuple(h.permute(1, 0, 2).contiguous() for hs in hidden_states for h in hs)
+ else:
+ hidden_states = tuple(hs.permute(1, 0, 2).contiguous() for hs in hidden_states)
+
+ if output_attentions:
+ if target_mapping is not None:
+ # when target_mapping is provided, there are 2-tuple of attentions
+ attentions = tuple(
+ tuple(att_stream.permute(2, 3, 0, 1).contiguous() for att_stream in t) for t in attentions
+ )
+ else:
+ attentions = tuple(t.permute(2, 3, 0, 1).contiguous() for t in attentions)
+
+ if not return_dict:
+ return tuple(v for v in [output, new_mems, hidden_states, attentions] if v is not None)
+
+ return XLNetModelOutput(
+ last_hidden_state=output, mems=new_mems, hidden_states=hidden_states, attentions=attentions
+ )
+
+
+@add_start_docstrings(
+ """
+ XLNet Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetLMHeadModel(XLNetPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = ["lm_loss.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.attn_type = config.attn_type
+ self.same_length = config.same_length
+
+ self.transformer = XLNetModel(config)
+ self.lm_loss = nn.Linear(config.d_model, config.vocab_size, bias=True)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.lm_loss
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_loss = new_embeddings
+
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, use_mems=None, **kwargs):
+ # Overwritten -- this model has unique input preparation
+
+ # Add dummy token at the end (no attention on this one)
+
+ effective_batch_size = input_ids.shape[0]
+ dummy_token = torch.zeros((effective_batch_size, 1), dtype=torch.long, device=input_ids.device)
+
+ # At every pass, the attention values for the new token and the two last generated tokens
+ # are computed, the rest is reloaded from the `past` cache. A purely auto-regressive model would have
+ # offset = 1; offset = 2 seems to have slightly better computation.
+ offset = 2
+
+ if past_key_values:
+ input_ids = torch.cat([input_ids[:, -offset:], dummy_token], dim=1)
+ else:
+ input_ids = torch.cat([input_ids, dummy_token], dim=1)
+
+ # Build permutation mask so that previous tokens don't see last token
+ sequence_length = input_ids.shape[1]
+ perm_mask = torch.zeros(
+ (effective_batch_size, sequence_length, sequence_length), dtype=torch.float, device=input_ids.device
+ )
+ perm_mask[:, :, -1] = 1.0
+
+ # We'll only predict the last token
+ target_mapping = torch.zeros(
+ (effective_batch_size, 1, sequence_length), dtype=torch.float, device=input_ids.device
+ )
+ target_mapping[:, 0, -1] = 1.0
+
+ inputs = {
+ "input_ids": input_ids,
+ "perm_mask": perm_mask,
+ "target_mapping": target_mapping,
+ "use_mems": use_mems,
+ }
+
+ # if past is defined in model kwargs then use it for faster decoding
+ if past_key_values:
+ inputs["mems"] = tuple(layer_past[:-offset, :, :] for layer_past in past_key_values)
+
+ return inputs
+
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=XLNetLMHeadModelOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_mask: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetLMHeadModelOutput]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, num_predict)`, *optional*):
+ Labels for masked language modeling. `num_predict` corresponds to `target_mapping.shape[1]`. If
+ `target_mapping` is `None`, then `num_predict` corresponds to `sequence_length`.
+
+ The labels should correspond to the masked input words that should be predicted and depends on
+ `target_mapping`. Note in order to perform standard auto-regressive language modeling a ** token has
+ to be added to the `input_ids` (see the `prepare_inputs_for_generation` function and examples below)
+
+ Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored, the loss
+ is only computed for labels in `[0, ..., config.vocab_size]`
+
+ Return:
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, XLNetLMHeadModel
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("xlnet/xlnet-large-cased")
+ >>> model = XLNetLMHeadModel.from_pretrained("xlnet/xlnet-large-cased")
+
+ >>> # We show how to setup inputs to predict a next token using a bi-directional context.
+ >>> input_ids = torch.tensor(
+ ... tokenizer.encode("Hello, my dog is very ", add_special_tokens=False)
+ ... ).unsqueeze(
+ ... 0
+ ... ) # We will predict the masked token
+ >>> perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float)
+ >>> perm_mask[:, :, -1] = 1.0 # Previous tokens don't see last token
+ >>> target_mapping = torch.zeros(
+ ... (1, 1, input_ids.shape[1]), dtype=torch.float
+ ... ) # Shape [1, 1, seq_length] => let's predict one token
+ >>> target_mapping[
+ ... 0, 0, -1
+ ... ] = 1.0 # Our first (and only) prediction will be the last token of the sequence (the masked token)
+
+ >>> outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping)
+ >>> next_token_logits = outputs[
+ ... 0
+ ... ] # Output has shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size]
+
+ >>> # The same way can the XLNetLMHeadModel be used to be trained by standard auto-regressive language modeling.
+ >>> input_ids = torch.tensor(
+ ... tokenizer.encode("Hello, my dog is very ", add_special_tokens=False)
+ ... ).unsqueeze(
+ ... 0
+ ... ) # We will predict the masked token
+ >>> labels = torch.tensor(tokenizer.encode("cute", add_special_tokens=False)).unsqueeze(0)
+ >>> assert labels.shape[0] == 1, "only one word will be predicted"
+ >>> perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float)
+ >>> perm_mask[
+ ... :, :, -1
+ ... ] = 1.0 # Previous tokens don't see last token as is done in standard auto-regressive lm training
+ >>> target_mapping = torch.zeros(
+ ... (1, 1, input_ids.shape[1]), dtype=torch.float
+ ... ) # Shape [1, 1, seq_length] => let's predict one token
+ >>> target_mapping[
+ ... 0, 0, -1
+ ... ] = 1.0 # Our first (and only) prediction will be the last token of the sequence (the masked token)
+
+ >>> outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping, labels=labels)
+ >>> loss = outputs.loss
+ >>> next_token_logits = (
+ ... outputs.logits
+ ... ) # Logits have shape [target_mapping.size(0), target_mapping.size(1), 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+
+ logits = self.lm_loss(transformer_outputs[0])
+
+ loss = None
+ if labels is not None:
+ # Flatten the tokens
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return XLNetLMHeadModelOutput(
+ loss=loss,
+ logits=logits,
+ mems=transformer_outputs.mems,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+ @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(
+ """
+ XLNet Model with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g.
+ for GLUE tasks.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetForSequenceClassification(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.config = config
+
+ self.transformer = XLNetModel(config)
+ self.sequence_summary = SequenceSummary(config)
+ self.logits_proj = nn.Linear(config.d_model, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=XLNetForSequenceClassificationOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_mask: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetForSequenceClassificationOutput]:
+ 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+ output = transformer_outputs[0]
+
+ output = self.sequence_summary(output)
+ logits = self.logits_proj(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 XLNetForSequenceClassificationOutput(
+ loss=loss,
+ logits=logits,
+ mems=transformer_outputs.mems,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ XLNet 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.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetForTokenClassification(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.transformer = XLNetModel(config)
+ 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(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=XLNetForTokenClassificationOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_mask: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetForTokenClassificationOutput]:
+ r"""
+ 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_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.transformer(
+ input_ids,
+ attention_mask=attention_mask,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ 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:
+ 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 XLNetForTokenClassificationOutput(
+ loss=loss,
+ logits=logits,
+ mems=outputs.mems,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ XLNet Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
+ softmax) e.g. for RACE/SWAG tasks.
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetForMultipleChoice(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.transformer = XLNetModel(config)
+ self.sequence_summary = SequenceSummary(config)
+ self.logits_proj = nn.Linear(config.d_model, 1)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=XLNetForMultipleChoiceOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_mask: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ labels: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetForMultipleChoiceOutput]:
+ 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_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_input_mask = input_mask.view(-1, input_mask.size(-1)) if input_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
+ )
+
+ transformer_outputs = self.transformer(
+ flat_input_ids,
+ token_type_ids=flat_token_type_ids,
+ input_mask=flat_input_mask,
+ attention_mask=flat_attention_mask,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ head_mask=head_mask,
+ inputs_embeds=flat_inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+
+ output = transformer_outputs[0]
+
+ output = self.sequence_summary(output)
+ logits = self.logits_proj(output)
+ reshaped_logits = logits.view(-1, num_choices)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(reshaped_logits, labels.view(-1))
+
+ if not return_dict:
+ output = (reshaped_logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return XLNetForMultipleChoiceOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ mems=transformer_outputs.mems,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ XLNet 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`).
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetForQuestionAnsweringSimple(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.transformer = XLNetModel(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(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @add_code_sample_docstrings(
+ checkpoint=_CHECKPOINT_FOR_DOC,
+ output_type=XLNetForQuestionAnsweringSimpleOutput,
+ config_class=_CONFIG_FOR_DOC,
+ )
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_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,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetForQuestionAnsweringSimpleOutput]:
+ 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+
+ 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 XLNetForQuestionAnsweringSimpleOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ mems=outputs.mems,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@add_start_docstrings(
+ """
+ XLNet 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`).
+ """,
+ XLNET_START_DOCSTRING,
+)
+class XLNetForQuestionAnswering(XLNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.start_n_top = config.start_n_top
+ self.end_n_top = config.end_n_top
+
+ self.transformer = XLNetModel(config)
+ self.start_logits = PoolerStartLogits(config)
+ self.end_logits = PoolerEndLogits(config)
+ self.answer_class = PoolerAnswerClass(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+ @replace_return_docstrings(output_type=XLNetForQuestionAnsweringOutput, config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ mems: Optional[torch.Tensor] = None,
+ perm_mask: Optional[torch.Tensor] = None,
+ target_mapping: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ input_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,
+ is_impossible: Optional[torch.Tensor] = None,
+ cls_index: Optional[torch.Tensor] = None,
+ p_mask: Optional[torch.Tensor] = None,
+ use_mems: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ **kwargs, # delete when `use_cache` is removed in XLNetModel
+ ) -> Union[Tuple, XLNetForQuestionAnsweringOutput]:
+ 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 AutoTokenizer, XLNetForQuestionAnswering
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("xlnet/xlnet-base-cased")
+ >>> model = XLNetForQuestionAnswering.from_pretrained("xlnet/xlnet-base-cased")
+
+ >>> 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,
+ mems=mems,
+ perm_mask=perm_mask,
+ target_mapping=target_mapping,
+ token_type_ids=token_type_ids,
+ input_mask=input_mask,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ use_mems=use_mems,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+ hidden_states = transformer_outputs[0]
+ start_logits = self.start_logits(hidden_states, p_mask=p_mask)
+
+ outputs = transformer_outputs[1:] # Keep mems, hidden states, attentions if there are in it
+
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, let's remove the dimension added by batch splitting
+ for x in (start_positions, end_positions, cls_index, is_impossible):
+ if x is not None and x.dim() > 1:
+ x.squeeze_(-1)
+
+ # during training, compute the end logits based on the ground truth of the start position
+ end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)
+
+ loss_fct = CrossEntropyLoss()
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ if cls_index is not None and is_impossible is not None:
+ # Predict answerability from the representation of CLS and START
+ cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
+ loss_fct_cls = nn.BCEWithLogitsLoss()
+ cls_loss = loss_fct_cls(cls_logits, is_impossible)
+
+ # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
+ total_loss += cls_loss * 0.5
+
+ if not return_dict:
+ return (total_loss,) + transformer_outputs[1:]
+ else:
+ return XLNetForQuestionAnsweringOutput(
+ loss=total_loss,
+ mems=transformer_outputs.mems,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+ else:
+ # during inference, compute the end logits based on beam search
+ bsz, slen, hsz = hidden_states.size()
+ start_log_probs = nn.functional.softmax(start_logits, dim=-1) # shape (bsz, slen)
+
+ start_top_log_probs, start_top_index = torch.topk(
+ start_log_probs, self.start_n_top, dim=-1
+ ) # shape (bsz, start_n_top)
+ start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz)
+ start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz)
+ start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz)
+
+ hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(
+ start_states
+ ) # shape (bsz, slen, start_n_top, hsz)
+ p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
+ end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
+ end_log_probs = nn.functional.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top)
+
+ end_top_log_probs, end_top_index = torch.topk(
+ end_log_probs, self.end_n_top, dim=1
+ ) # shape (bsz, end_n_top, start_n_top)
+ end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
+ end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)
+
+ start_states = torch.einsum(
+ "blh,bl->bh", hidden_states, start_log_probs
+ ) # get the representation of START as weighted sum of hidden states
+ cls_logits = self.answer_class(
+ hidden_states, start_states=start_states, cls_index=cls_index
+ ) # Shape (batch size,): one single `cls_logits` for each sample
+
+ if not return_dict:
+ outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits)
+ return outputs + transformer_outputs[1:]
+ else:
+ return XLNetForQuestionAnsweringOutput(
+ start_top_log_probs=start_top_log_probs,
+ start_top_index=start_top_index,
+ end_top_log_probs=end_top_log_probs,
+ end_top_index=end_top_index,
+ cls_logits=cls_logits,
+ mems=transformer_outputs.mems,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+__all__ = [
+ "XLNetForMultipleChoice",
+ "XLNetForQuestionAnswering",
+ "XLNetForQuestionAnsweringSimple",
+ "XLNetForSequenceClassification",
+ "XLNetForTokenClassification",
+ "XLNetLMHeadModel",
+ "XLNetModel",
+ "XLNetPreTrainedModel",
+ "load_tf_weights_in_xlnet",
+]
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet.py b/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..640558e1d81f7d2109e02bc418a29e1b3d0e2c78
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet.py
@@ -0,0 +1,385 @@
+# coding=utf-8
+# Copyright 2018 Google AI, 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.
+"""Tokenization classes for XLNet 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 SPIECE_UNDERLINE, logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
+
+
+# Segments (not really needed)
+SEG_ID_A = 0
+SEG_ID_B = 1
+SEG_ID_CLS = 2
+SEG_ID_SEP = 3
+SEG_ID_PAD = 4
+
+
+class XLNetTokenizer(PreTrainedTokenizer):
+ """
+ Construct an XLNet 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.
+
+ 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 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.
+ 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
+ padding_side = "left"
+
+ 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:
+ # Mask token behave like a normal word, i.e. include the space before it
+ mask_token = AddedToken(mask_token, lstrip=True, 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,
+ 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
+ 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)
+
+ 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 _decode(
+ self,
+ token_ids: List[int],
+ skip_special_tokens: bool = False,
+ clean_up_tokenization_spaces: Optional[bool] = None,
+ spaces_between_special_tokens: bool = True,
+ **kwargs,
+ ) -> str:
+ self._decode_use_source_tokenizer = kwargs.pop("use_source_tokenizer", False)
+
+ filtered_tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens)
+
+ # To avoid mixing byte-level and unicode for byte-level BPT
+ # we need to build string separately for added tokens and byte-level tokens
+ # cf. https://github.com/huggingface/transformers/issues/1133
+ sub_texts = []
+ current_sub_text = []
+ for token in filtered_tokens:
+ if skip_special_tokens and token in self.all_special_ids:
+ continue
+ if token in self.added_tokens_encoder:
+ if current_sub_text:
+ sub_texts.append(self.convert_tokens_to_string(current_sub_text))
+ current_sub_text = []
+ sub_texts.append(token)
+ else:
+ current_sub_text.append(token)
+ if current_sub_text:
+ sub_texts.append(self.convert_tokens_to_string(current_sub_text))
+
+ # Mimic the behavior of the Rust tokenizer:
+ # By default, there are no spaces between special tokens
+ text = "".join(sub_texts)
+
+ clean_up_tokenization_spaces = (
+ clean_up_tokenization_spaces
+ if clean_up_tokenization_spaces is not None
+ else self.clean_up_tokenization_spaces
+ )
+ if clean_up_tokenization_spaces:
+ clean_text = self.clean_up_tokenization(text)
+ return clean_text
+ else:
+ 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 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
+
+ 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]
+
+ 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
+
+ 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,)
+
+
+__all__ = ["XLNetTokenizer"]
diff --git a/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet_fast.py b/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a6ae907ef573f16cb46ad7eae335ac6894e0c37
--- /dev/null
+++ b/lib/python3.12/site-packages/transformers/models/xlnet/tokenization_xlnet_fast.py
@@ -0,0 +1,234 @@
+# coding=utf-8
+# Copyright 2018 Google AI, 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.
+"""Tokenization classes for XLNet 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_xlnet import XLNetTokenizer
+else:
+ XLNetTokenizer = None
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}
+
+
+SPIECE_UNDERLINE = "▁"
+
+# Segments (not really needed)
+SEG_ID_A = 0
+SEG_ID_B = 1
+SEG_ID_CLS = 2
+SEG_ID_SEP = 3
+SEG_ID_PAD = 4
+
+
+class XLNetTokenizerFast(PreTrainedTokenizerFast):
+ """
+ Construct a "fast" XLNet tokenizer (backed by HuggingFace's *tokenizers* library). 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 `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).
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ padding_side = "left"
+ slow_tokenizer_class = XLNetTokenizer
+
+ 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,
+ ):
+ # 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
+
+ @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 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
+
+ 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
+
+ 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,)
+
+
+__all__ = ["XLNetTokenizerFast"]
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1053251ed7c2ad2ba3ce8bd18d184980e2eb34d2
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/attention_visualizer.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/attention_visualizer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bea50e784efee5c73e15d3f33e790be3c7cdb62d
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/attention_visualizer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/backbone_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/backbone_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ad3caa6cb4166aef94738864a874b214669717a6
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/backbone_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/bitsandbytes.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/bitsandbytes.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a6a4b9db3cbad8ca90385282c2b40b8c83e57b12
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/bitsandbytes.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/deprecation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/deprecation.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f3da7ec2633c2e3aaaaa9d6df4eb0bfb3e4ead15
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/deprecation.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/doc.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/doc.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f9f01fe789070ad52accdffb88518ddf79541273
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/doc.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_detectron2_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_detectron2_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f5d90e70ba0c16fd5ae3acffa24961b6e68c55b4
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_detectron2_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_flax_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_flax_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8e009cfe900c89bd545d6d03ea7e7f9593850224
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_flax_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_keras_nlp_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_keras_nlp_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..df4db94e0e59d3d4c672c0c787bf97aa2cd2f907
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_keras_nlp_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_music_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_music_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1137b62a87cb86618d1466a13efe0874274da1bb
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_music_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_and_tokenizers_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_and_tokenizers_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2b6da10a20ffed3751f75d57d19816f9845090fe
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_and_tokenizers_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..281f87a426d14487809e77fa8ac90b70f60377e7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_sentencepiece_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tensorflow_text_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tensorflow_text_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f482456f1455d8ffdfa2098bc75337917a0a8bc0
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tensorflow_text_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_timm_and_torchvision_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_timm_and_torchvision_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e80192d9e26c405700937d42f4130c51a89b4f2f
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_timm_and_torchvision_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tokenizers_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tokenizers_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..00c606ae0a384b66f36552c778202570ce7070b4
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_tokenizers_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchaudio_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchaudio_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..73f7d399db6a811f8787a6211e7d647a5fccac67
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchaudio_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchvision_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchvision_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9c1f3686735e04d5cd9695af2fff1840a80352a5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_torchvision_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_vision_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_vision_objects.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..03f781e215fea79a90b60dcf0841e48fc24c9c8f
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_vision_objects.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/fx.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/fx.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bf93c4752531c860b4d5fdfd28d498cd0bb4ce14
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/fx.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/generic.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/generic.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..aa4a347441f10e2fa51e437eabf0cbc01b2bafbe
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/generic.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/hub.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/hub.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..31a7e243b649415d2af9983a2b550173c36386b5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/hub.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/import_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/import_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8bb2d4b9be31b47c853740b2fb714b93ad5672ed
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/import_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/logging.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..177b5a7827c3494ca7f670fc275e5280e5026ed7
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/logging.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/model_parallel_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/model_parallel_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..691aaeaf118b2732e5fe646d6d0c4859868aa06c
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/model_parallel_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/notebook.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/notebook.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2ce5ea58510ae423e31b779545bb2c4a59f287b5
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/notebook.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/peft_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/peft_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bd149700a2d66327efa74970ad07cda5284e5ea3
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/peft_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/quantization_config.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/quantization_config.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..89f1d569d685d560bcd7137f2d69c5241f3c313a
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/quantization_config.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9b6d06592d0ca41af4e5e960d55ba2e642a51d67
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2_new.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2_new.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..55cf59211f2ab05e7e9b723dffd5d306031aa5ba
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/sentencepiece_model_pb2_new.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/versions.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/versions.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0f955c7594785a34811b60de9b2901369fee6bfb
Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/versions.cpython-312.pyc differ