model: Add GTE to Transformers

@@ -705,6 +705,8 @@         title: GraniteMoeSWA       - local: model_doc/granite_swa         title: GraniteSWA+      - local: model_doc/gte+        title: GTE       - local: model_doc/helium         title: Helium       - local: model_doc/herbert

@@ -0,0 +1,106 @@+<!--Copyright 2026 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.+++⚠️ Note that this file is in Markdown but contains specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer.++-->+*This model was published in HF papers on 2024-07-29 and contributed to Hugging Face Transformers on 2026-09-30.*++# GTE++## Overview++GTE was proposed in [mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval](https://huggingface.co/papers/2407.19669) by Xin Zhang, Yanzhao Zhang, Dingkun Long, Wen Xie, Ziqi Dai, Jialong Tang, Huan Lin, Baosong Yang, Pengjun Xie, Fei Huang, Meishan Zhang, Wenjie Li and Min Zhang.++The abstract from the paper is the following:++*We present systematic efforts in building long-context multilingual text representation model (TRM) and reranker from scratch for text retrieval. We first introduce a text encoder (base size) enhanced with RoPE and unpadding, pre-trained in a native 8192-token context (longer than 512 of previous multilingual encoders). Then we construct a hybrid TRM and a cross-encoder reranker by contrastive learning. Evaluations show that our text encoder outperforms the same-sized previous state-of-the-art XLM-R. Meanwhile, our TRM and reranker match the performance of large-sized state-of-the-art BGE-M3 models and achieve better results on long-context retrieval benchmarks. Further analysis demonstrate that our proposed models exhibit higher efficiency during both training and inference. We believe their efficiency and effectiveness could benefit various researches and industrial applications.*++GTE is a BERT-style bidirectional encoder that replaces absolute position embeddings with RoPE, uses a gated MLP, and applies layer normalization after each residual connection. The same architecture backs Alibaba's `gte-*-v1.5`, `gte-multilingual-*` and `gte-en-mlm-*` checkpoints as well as Snowflake's `snowflake-arctic-embed-m-v2.0`.++This model was contributed by [Harshal Janjani](https://huggingface.co/harshaljanjani).+The original code can be found [here](https://huggingface.co/Alibaba-NLP/new-impl).++> [!TIP]+> Click on the GTE models in the right sidebar for more examples of how to apply GTE to different language tasks.++The example below demonstrates how to extract features (embeddings) with [`Pipeline`] and [`AutoModel`].++<hfoptions id="usage">+<hfoption id="Pipeline">++```python+from transformers import pipeline+++pipeline = pipeline(+    task="feature-extraction",+    model="alibaba-nlp-community/gte-multilingual-base",+)+pipeline("Plants create oxygen through a process known as photosynthesis.")+```++</hfoption>+<hfoption id="AutoModel">++```python+import torch++from transformers import AutoModel, AutoTokenizer+++tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")+model = AutoModel.from_pretrained("alibaba-nlp-community/gte-multilingual-base", device_map="auto")+inputs = tokenizer("Plants create oxygen through a process known as photosynthesis.", return_tensors="pt").to(model.device)++with torch.no_grad():+    outputs = model(**inputs)+    embeddings = outputs.last_hidden_state[:, 0]++print(f"Embeddings shape: {embeddings.shape}")+```++</hfoption>+</hfoptions>++## Notes++- GTE uses RoPE, so for correct positional encoding either use right padding (the default), or use left padding and prepare `position_ids` accordingly.+- `type_vocab_size` differs across checkpoints. `alibaba-nlp-community/gte-base-en-v1.5` sets it to `0`, in which case no token type embedding is created and `token_type_ids` are ignored.+- The `gte-*-v1.5` and `gte-multilingual-*` checkpoints apply static NTK scaling on top of RoPE. It is expressed as a `linear` [`~modeling_rope_utils.RopeParameters`] entry whose `rope_theta` is the base scaled by the NTK factor.++## GteConfig++[[autodoc]] GteConfig++## GteModel++[[autodoc]] GteModel+    - forward++## GteForMaskedLM++[[autodoc]] GteForMaskedLM+    - forward++## GteForSequenceClassification++[[autodoc]] GteForSequenceClassification+    - forward++## GteForTokenClassification++[[autodoc]] GteForTokenClassification+    - forward

@@ -1407,6 +1407,31 @@ def _build_checkpoint_conversion_mapping():                 operations=[Chunk(dim=0)],             ),         ],+        "gte": [+            PrefixChange(prefix_to_remove="new"),+            WeightRenaming(r"encoder.layer", r"layers"),+            WeightRenaming(r"attention.o_proj", r"self_attn.o_proj"),+            WeightRenaming(r"attn_ln", r"post_attention_layernorm"),+            WeightRenaming(r"mlp_ln", r"post_mlp_layernorm"),+            WeightRenaming(r"lm_head.norm", r"lm_head.layer_norm"),+            WeightConverter(+                source_patterns="attention.qkv_proj",+                target_patterns=[+                    "self_attn.q_proj",+                    "self_attn.k_proj",+                    "self_attn.v_proj",+                ],+                operations=[Chunk(dim=0)],+            ),+            WeightConverter(+                source_patterns="mlp.up_gate_proj",+                target_patterns=[+                    "mlp.up_proj",+                    "mlp.gate_proj",+                ],+                operations=[Chunk(dim=0)],+            ),+        ],         "jina_embeddings_v3": [             WeightRenaming(source_patterns="emb_ln", target_patterns="embeddings.LayerNorm"),             WeightRenaming(source_patterns="encoder.layers", target_patterns="layers"),

@@ -212,6 +212,7 @@     from .granitemoeshared import *     from .grounding_dino import *     from .groupvit import *+    from .gte import *     from .helium import *     from .herbert import *     from .hgnet_v2 import *

@@ -288,6 +288,7 @@         ("groupvit", "GroupViTConfig"),         ("groupvit_text_model", "GroupViTTextConfig"),         ("groupvit_vision_model", "GroupViTVisionConfig"),+        ("gte", "GteConfig"),         ("helium", "HeliumConfig"),         ("hgnet_v2", "HGNetV2Config"),         ("hiera", "HieraConfig"),

@@ -248,6 +248,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("granitemoeshared", "GraniteMoeSharedModel"),         ("grounding-dino", "GroundingDinoModel"),         ("groupvit", "GroupViTModel"),+        ("gte", "GteModel"),         ("helium", "HeliumModel"),         ("hgnet_v2", "HGNetV2Backbone"),         ("hiera", "HieraModel"),@@ -1230,6 +1231,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("flaubert", "FlaubertWithLMHeadModel"),         ("fnet", "FNetForMaskedLM"),         ("funnel", "FunnelForMaskedLM"),+        ("gte", "GteForMaskedLM"),         ("ibert", "IBertForMaskedLM"),         ("jina_embeddings_v3", "JinaEmbeddingsV3ForMaskedLM"),         ("layoutlm", "LayoutLMForMaskedLM"),@@ -1448,6 +1450,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("gpt_neox", "GPTNeoXForSequenceClassification"),         ("gpt_oss", "GptOssForSequenceClassification"),         ("gptj", "GPTJForSequenceClassification"),+        ("gte", "GteForSequenceClassification"),         ("helium", "HeliumForSequenceClassification"),         ("hunyuan_v1_dense", "HunYuanDenseV1ForSequenceClassification"),         ("hunyuan_v1_moe", "HunYuanMoEV1ForSequenceClassification"),@@ -1677,6 +1680,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("gpt_neo", "GPTNeoForTokenClassification"),         ("gpt_neox", "GPTNeoXForTokenClassification"),         ("gpt_oss", "GptOssForTokenClassification"),+        ("gte", "GteForTokenClassification"),         ("helium", "HeliumForTokenClassification"),         ("ibert", "IBertForTokenClassification"),         ("jina_embeddings_v3", "JinaEmbeddingsV3ForTokenClassification"),

@@ -161,6 +161,7 @@         ("granitemoeshared", "TokenizersBackend" if is_tokenizers_available() else None),         ("grounding-dino", "BertTokenizer" if is_tokenizers_available() else None),         ("groupvit", "CLIPTokenizer" if is_tokenizers_available() else None),+        ("gte", "XLMRobertaTokenizer" if is_tokenizers_available() else None),         ("herbert", "HerbertTokenizer" if is_tokenizers_available() else None),         ("hubert", "Wav2Vec2CTCTokenizer"),         ("hunyuan_vl", "Qwen2Tokenizer" if is_tokenizers_available() else None),

@@ -0,0 +1,33 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from the modules in `src/transformers/models/gte`.+#           Do NOT edit this file manually as any edits will be overwritten by auto-generation of the file.+#                       A module is picked up once it exports a name through `__all__`.+#                   Regenerate the file with: `python utils/check_inits.py --fix_and_overwrite`+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 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_gte import *+    from .modeling_gte import *+else:+    import sys++    _file = globals()["__file__"]+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

@@ -0,0 +1,70 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from src/transformers/models/gte/modular_gte.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_gte.py file directly. One of our CI enforces this.+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 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 huggingface_hub.dataclasses import strict++from ...configuration_utils import PreTrainedConfig+from ...modeling_rope_utils import RopeParameters+from ...utils import auto_docstring+++@auto_docstring(checkpoint="alibaba-nlp-community/gte-multilingual-base")+@strict+class GteConfig(PreTrainedConfig):+    r"""+    Examples:++    ```python+    >>> from transformers import GteConfig, GteModel++    >>> # Initializing a GTE alibaba-nlp-community/gte-multilingual-base style configuration+    >>> configuration = GteConfig()++    >>> # Initializing a model (with random weights) from the alibaba-nlp-community/gte-multilingual-base style configuration+    >>> model = GteModel(configuration)++    >>> # Accessing the model configuration+    >>> configuration = model.config+    ```"""++    model_type = "gte"++    vocab_size: int = 250048+    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 | int = 0.1+    attention_probs_dropout_prob: float | int = 0.0+    max_position_embeddings: int = 8192+    type_vocab_size: int = 1+    initializer_range: float = 0.02+    layer_norm_eps: float = 1e-12+    pad_token_id: int | None = 1+    classifier_dropout: float | int | None = None+    bos_token_id: int | None = None+    eos_token_id: int | list[int] | None = None+    tie_word_embeddings: bool = True+    default_theta = 160_000.0+    rope_parameters: RopeParameters | dict | None = None+++__all__ = ["GteConfig"]

@@ -0,0 +1,660 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from src/transformers/models/gte/modular_gte.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_gte.py file directly. One of our CI enforces this.+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 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 collections.abc import Callable++import torch+import torch.nn as nn+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss++from ...activations import ACT2FN, gelu+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func+from ...masking_utils import create_bidirectional_mask+from ...modeling_layers import GradientCheckpointingLayer+from ...modeling_outputs import (+    BaseModelOutputWithPooling,+    MaskedLMOutput,+    SequenceClassifierOutput,+    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 TransformersKwargs, auto_docstring+from ...utils.deprecation import deprecate_kwarg+from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults+from ...utils.output_capturing import capture_outputs+from .configuration_gte import GteConfig+++class GteEmbeddings(nn.Module):+    """Construct the embeddings from word, position and token_type embeddings."""++    def __init__(self, config: GteConfig):+        super().__init__()+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)+        self.token_type_embeddings = (+            # CODEPATH: Only gte-base-en-v1.5 sets type_vocab_size=0, the others ship token type embeddings.+            nn.Embedding(config.type_vocab_size, config.hidden_size) if config.type_vocab_size else None+        )++        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)+        self.dropout = nn.Dropout(config.hidden_dropout_prob)++    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        token_type_ids: torch.LongTensor | None = None,+        position_ids: torch.LongTensor | None = None,+        inputs_embeds: torch.FloatTensor | None = None,+    ) -> torch.Tensor:+        # Unlike JinaEmbeddingsV3Embeddings, the token type embedding is skipped when type_vocab_size is 0.+        embeddings = inputs_embeds+        if inputs_embeds is None:+            embeddings = self.word_embeddings(input_ids)++        input_shape = embeddings.shape[:-1]+        device = embeddings.device++        if self.token_type_embeddings is not None:+            if token_type_ids is None:+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)+            embeddings = embeddings + self.token_type_embeddings(token_type_ids)++        embeddings = self.LayerNorm(embeddings)+        embeddings = self.dropout(embeddings)++        return embeddings+++class GteRotaryEmbedding(nn.Module):+    @deprecate_kwarg("device", version="5.18")+    def __init__(self, config: GteConfig, device=None):+        super().__init__()+        self.max_seq_len_cached = config.max_position_embeddings+        self.original_max_seq_len = config.max_position_embeddings++        self.config = config++        self.rope_type = self.config.rope_parameters["rope_type"]+        rope_init_fn: Callable = self.compute_default_rope_parameters+        if self.rope_type != "default":+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)++        self.inv_freq = nn.Buffer(inv_freq, persistent=False)+        self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)++    @staticmethod+    @deprecate_kwarg("device", version="5.18")+    def compute_default_rope_parameters(config: GteConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]:+        """+        Computes the inverse frequencies according to the original RoPE implementation+        Args:+            config ([`~transformers.PreTrainedConfig`]):+                The model configuration.+        Returns:+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).+        """+        base = config.rope_parameters["rope_theta"]+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads++        attention_factor = 1.0  # Unused in this type of RoPE+        # Compute the inverse frequencies+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))+        return inv_freq.to(device), attention_factor++    @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].expand(position_ids.shape[0], -1, 1).to(dtype=torch.float, device=x.device)+        )+        position_ids_expanded = position_ids[:, None, :].float()++        device_type = x.device.type if isinstance(x.device.type, str) else "cpu"+        # Disable any outside autocast context if any, to really force fp32+        with maybe_autocast(device_type=device_type, enabled=False):+            freqs = (inv_freq_expanded @ position_ids_expanded).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)+++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)+++@use_kernel_forward_from_hub("rotary_pos_emb")+def apply_rotary_pos_emb(q, k, cos, sin, 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.+        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 eager_attention_forward(+    module: nn.Module,+    query: torch.Tensor,+    key: torch.Tensor,+    value: torch.Tensor,+    attention_mask: torch.Tensor | None,+    scaling: float | None = None,+    dropout: float = 0.0,+    **kwargs: Unpack[TransformersKwargs],+):+    if scaling is None:+        scaling = query.size(-1) ** -0.5++    # Take the dot product between "query" and "key" to get the raw attention scores.+    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling++    if attention_mask is not None:+        attn_weights = attn_weights + attention_mask++    attn_weights = nn.functional.softmax(attn_weights, dim=-1)+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)++    attn_output = torch.matmul(attn_weights, value)+    attn_output = attn_output.transpose(1, 2).contiguous()++    return attn_output, attn_weights+++@use_kernelized_func(apply_rotary_pos_emb)+class GteAttention(nn.Module):+    """Multi-headed attention from 'Attention Is All You Need' paper"""++    def __init__(self, config: GteConfig):+        super().__init__()+        self.config = config+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)+        self.scaling = self.head_dim**-0.5+        self.attention_dropout = config.attention_probs_dropout_prob+        self.is_causal = False++        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)+        self.k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)+        self.v_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True)++    def forward(+        self,+        hidden_states: torch.Tensor,+        attention_mask: torch.Tensor | None = None,+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> 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)++        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(+            self.config._attn_implementation, eager_attention_forward+        )++        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 GteMLP(nn.Module):+    def __init__(self, config: GteConfig):+        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(config.intermediate_size, config.hidden_size, bias=True)+        self.act_fn = ACT2FN[config.hidden_act]+        self.dropout = nn.Dropout(config.hidden_dropout_prob)++    def forward(self, x: torch.Tensor) -> torch.Tensor:+        # Unlike GemmaMLP, `down_proj` carries a bias and the gated activation is dropped out.+        hidden_states = self.act_fn(self.gate_proj(x)) * self.up_proj(x)+        hidden_states = self.dropout(hidden_states)+        down_proj = self.down_proj(hidden_states)+        return down_proj+++class GteLayer(GradientCheckpointingLayer):+    def __init__(self, config: GteConfig):+        super().__init__()+        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)++        self.post_attention_dropout = nn.Dropout(config.hidden_dropout_prob)+        self.post_mlp_dropout = nn.Dropout(config.hidden_dropout_prob)+        self.mlp = GteMLP(config)+        self.self_attn = GteAttention(config=config)+        self.post_mlp_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)++    def forward(+        self,+        hidden_states: torch.Tensor,+        attention_mask: torch.Tensor | None = None,+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> torch.FloatTensor:+        residual = hidden_states+        attention_output, _ = self.self_attn(+            hidden_states,+            attention_mask=attention_mask,+            position_embeddings=position_embeddings,+            **kwargs,+        )+        hidden_states = residual + self.post_attention_dropout(attention_output)+        hidden_states = self.post_attention_layernorm(hidden_states)++        residual = hidden_states+        hidden_states = self.mlp(hidden_states)+        hidden_states = residual + self.post_mlp_dropout(hidden_states)+        hidden_states = self.post_mlp_layernorm(hidden_states)+        return hidden_states+++class GtePooler(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+++@auto_docstring+class GtePreTrainedModel(PreTrainedModel):+    config_class = GteConfig+    base_model_prefix = "gte"+    supports_gradient_checkpointing = True+    _supports_flash_attn = True+    _supports_sdpa = True+    _supports_flex_attn = True+    _supports_attention_backend = True+    _can_record_outputs = {+        "hidden_states": GteLayer,+        "attentions": GteAttention,+    }++    # Are kept as non-persistent buffers to avoid being saved in the state dict+    # and causing mismatch when loading from a checkpoint that doesn't have them+    _keys_to_ignore_on_load_unexpected = ["inv_freq", "original_inv_freq"]+++@auto_docstring+class GteModel(GtePreTrainedModel):+    _no_split_modules = ["GteEmbeddings", "GteLayer"]++    def __init__(self, config: GteConfig, add_pooling_layer=False):+        r"""+        add_pooling_layer (`bool`, *optional*, defaults to `False`):+            Whether to add a pooling layer.+        """+        super().__init__(config)+        self.config = config+        self.gradient_checkpointing = False++        self.embeddings = GteEmbeddings(config)++        self.pooler = GtePooler(config) if add_pooling_layer else None+        self.rotary_emb = GteRotaryEmbedding(config)+        self.layers = nn.ModuleList([GteLayer(config) for _ in range(config.num_hidden_layers)])++        # 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++    @merge_with_config_defaults+    @capture_outputs+    @auto_docstring+    def forward(+        self,+        input_ids: torch.Tensor | None = None,+        attention_mask: torch.Tensor | None = None,+        token_type_ids: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> BaseModelOutputWithPooling | tuple:+        if (input_ids is None) ^ (inputs_embeds is not None):+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")++        if input_ids is not None:+            seq_length = input_ids.shape[1]+            device = input_ids.device+        else:+            seq_length = inputs_embeds.shape[1]+            device = inputs_embeds.device++        if position_ids is None:+            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)[None, :]++        embedding_output = self.embeddings(+            input_ids=input_ids,+            position_ids=position_ids,+            token_type_ids=token_type_ids,+            inputs_embeds=inputs_embeds,+        )++        attention_mask = create_bidirectional_mask(+            config=self.config,+            inputs_embeds=embedding_output,+            attention_mask=attention_mask,+        )++        hidden_states = embedding_output+        position_embeddings = self.rotary_emb(embedding_output, position_ids)++        for encoder_layer in self.layers[: self.config.num_hidden_layers]:+            hidden_states = encoder_layer(+                hidden_states,+                attention_mask=attention_mask,+                position_embeddings=position_embeddings,+                **kwargs,+            )++        sequence_output = hidden_states+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None++        return BaseModelOutputWithPooling(+            last_hidden_state=sequence_output,+            pooler_output=pooled_output,+        )+++class GteLMHead(nn.Module):+    """Gte Head for masked language modeling."""++    def __init__(self, config: GteConfig):+        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)++    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+++@auto_docstring+class GteForMaskedLM(GtePreTrainedModel):+    _tied_weights_keys = {"lm_head.decoder.weight": "gte.embeddings.word_embeddings.weight"}++    def __init__(self, config: GteConfig):+        super().__init__(config)++        self.gte = GteModel(config, add_pooling_layer=False)+        self.lm_head = GteLMHead(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++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.FloatTensor | None = None,+        token_type_ids: torch.LongTensor | None = None,+        position_ids: torch.LongTensor | None = None,+        inputs_embeds: torch.FloatTensor | None = None,+        labels: torch.LongTensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | MaskedLMOutput:+        # Unlike JinaEmbeddingsV3ForMaskedLM, the backbone is bound as `self.gte`, which modular cannot rename.+        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]`+        """+        outputs = self.gte(+            input_ids=input_ids,+            attention_mask=attention_mask,+            token_type_ids=token_type_ids,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            return_dict=True,+            **kwargs,+        )+        sequence_output = outputs[0]+        prediction_scores = self.lm_head(sequence_output)++        masked_lm_loss = None+        if labels is not None:+            # move labels to correct device+            labels = labels.to(prediction_scores.device)+            loss_fct = CrossEntropyLoss()  # -100 index = padding token+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))++        return MaskedLMOutput(+            loss=masked_lm_loss,+            logits=prediction_scores,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++@auto_docstring(+    custom_intro="""+    Gte Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled+    output) e.g. for GLUE tasks.+    """+)+class GteForSequenceClassification(GtePreTrainedModel):+    def __init__(self, config: GteConfig):+        super().__init__(config)+        self.num_labels = config.num_labels+        self.config = config+        self.gte = GteModel(config, add_pooling_layer=True)+        classifier_dropout = (+            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob+        )+        self.dropout = nn.Dropout(classifier_dropout)+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)++        # Initialize weights and apply final processing+        self.post_init()++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.Tensor | None = None,+        attention_mask: torch.Tensor | None = None,+        token_type_ids: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        labels: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | SequenceClassifierOutput:+        outputs = self.gte(+            input_ids,+            attention_mask=attention_mask,+            token_type_ids=token_type_ids,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            return_dict=True,+            **kwargs,+        )++        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)++        return SequenceClassifierOutput(+            loss=loss,+            logits=logits,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++@auto_docstring+class GteForTokenClassification(GtePreTrainedModel):+    def __init__(self, config):+        super().__init__(config)+        self.num_labels = config.num_labels++        self.gte = GteModel(config, add_pooling_layer=False)+        classifier_dropout = (+            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob+        )+        self.dropout = nn.Dropout(classifier_dropout)+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)++        # Initialize weights and apply final processing+        self.post_init()++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.Tensor | None = None,+        attention_mask: torch.Tensor | None = None,+        token_type_ids: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        labels: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | TokenClassifierOutput:+        outputs = self.gte(+            input_ids,+            attention_mask=attention_mask,+            token_type_ids=token_type_ids,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            return_dict=True,+            **kwargs,+        )++        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))++        return TokenClassifierOutput(+            loss=loss,+            logits=logits,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++__all__ = [+    "GtePreTrainedModel",+    "GteModel",+    "GteForMaskedLM",+    "GteForSequenceClassification",+    "GteForTokenClassification",+]

@@ -0,0 +1,255 @@+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++import torch+import torch.nn as nn+from huggingface_hub.dataclasses import strict+from torch.nn import CrossEntropyLoss++from ...modeling_outputs import MaskedLMOutput+from ...modeling_rope_utils import RopeParameters+from ...modeling_utils import PreTrainedModel+from ...processing_utils import Unpack+from ...utils import TransformersKwargs, auto_docstring+from ...utils.generic import can_return_tuple, no_inherit_decorator+from ..bert.configuration_bert import BertConfig+from ..bert.modeling_bert import (+    BertForSequenceClassification,+    BertForTokenClassification,+    BertPreTrainedModel,+)+from ..gemma.modeling_gemma import GemmaMLP+from ..jina_embeddings_v3.modeling_jina_embeddings_v3 import (+    JinaEmbeddingsV3Attention,+    JinaEmbeddingsV3Embeddings,+    JinaEmbeddingsV3ForMaskedLM,+    JinaEmbeddingsV3Layer,+    JinaEmbeddingsV3LMHead,+    JinaEmbeddingsV3Model,+    JinaEmbeddingsV3Pooler,+)+from ..llama.modeling_llama import LlamaRotaryEmbedding+++@auto_docstring(checkpoint="alibaba-nlp-community/gte-multilingual-base")+@strict+class GteConfig(BertConfig):+    r"""+    Examples:++    ```python+    >>> from transformers import GteConfig, GteModel++    >>> # Initializing a GTE alibaba-nlp-community/gte-multilingual-base style configuration+    >>> configuration = GteConfig()++    >>> # Initializing a model (with random weights) from the alibaba-nlp-community/gte-multilingual-base style configuration+    >>> model = GteModel(configuration)++    >>> # Accessing the model configuration+    >>> configuration = model.config+    ```"""++    model_type = "gte"+    default_theta = 160_000.0++    vocab_size: int = 250048+    attention_probs_dropout_prob: float | int = 0.0+    type_vocab_size: int = 1+    pad_token_id: int | None = 1+    max_position_embeddings: int = 8192+    rope_parameters: RopeParameters | dict | None = None+    is_decoder = AttributeError()+    add_cross_attention = AttributeError()+    use_cache = AttributeError()+++class GteEmbeddings(JinaEmbeddingsV3Embeddings):+    def __init__(self, config: GteConfig):+        super().__init__(config)+        # GTE is rope-only: absolute position ids are never read, and token types always default to zeros.+        del self.position_ids+        del self.token_type_ids+        self.token_type_embeddings = (+            # CODEPATH: Only gte-base-en-v1.5 sets type_vocab_size=0, the others ship token type embeddings.+            nn.Embedding(config.type_vocab_size, config.hidden_size) if config.type_vocab_size else None+        )++    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        token_type_ids: torch.LongTensor | None = None,+        position_ids: torch.LongTensor | None = None,+        inputs_embeds: torch.FloatTensor | None = None,+    ) -> torch.Tensor:+        # Unlike JinaEmbeddingsV3Embeddings, the token type embedding is skipped when type_vocab_size is 0.+        embeddings = inputs_embeds+        if inputs_embeds is None:+            embeddings = self.word_embeddings(input_ids)++        input_shape = embeddings.shape[:-1]+        device = embeddings.device++        if self.token_type_embeddings is not None:+            if token_type_ids is None:+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)+            embeddings = embeddings + self.token_type_embeddings(token_type_ids)++        embeddings = self.LayerNorm(embeddings)+        embeddings = self.dropout(embeddings)++        return embeddings+++class GteRotaryEmbedding(LlamaRotaryEmbedding):+    pass+++class GteAttention(JinaEmbeddingsV3Attention):+    pass+++@no_inherit_decorator+class GteMLP(GemmaMLP):+    def __init__(self, config: GteConfig):+        super().__init__(config)+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=True)+        self.dropout = nn.Dropout(config.hidden_dropout_prob)++    def forward(self, x: torch.Tensor) -> torch.Tensor:+        # Unlike GemmaMLP, `down_proj` carries a bias and the gated activation is dropped out.+        hidden_states = self.act_fn(self.gate_proj(x)) * self.up_proj(x)+        hidden_states = self.dropout(hidden_states)+        down_proj = self.down_proj(hidden_states)+        return down_proj+++class GteLayer(JinaEmbeddingsV3Layer):+    pass+++class GtePooler(JinaEmbeddingsV3Pooler):+    pass+++class GtePreTrainedModel(BertPreTrainedModel):+    config_class = GteConfig+    base_model_prefix = "gte"++    # Are kept as non-persistent buffers to avoid being saved in the state dict+    # and causing mismatch when loading from a checkpoint that doesn't have them+    _keys_to_ignore_on_load_unexpected = ["inv_freq", "original_inv_freq"]+    _can_record_outputs = {+        "hidden_states": GteLayer,+        "attentions": GteAttention,+    }++    def _init_weights(self, **super_kwargs):+        raise AttributeError("Uses base super call")+++@auto_docstring+class GteModel(JinaEmbeddingsV3Model):+    def __init__(self, config: GteConfig, add_pooling_layer=False):+        r"""+        add_pooling_layer (`bool`, *optional*, defaults to `False`):+            Whether to add a pooling layer.+        """+        super().__init__(config, add_pooling_layer=add_pooling_layer)+++class GteLMHead(JinaEmbeddingsV3LMHead):+    def __init__(self, config: GteConfig):+        super().__init__(config)+        # GTE folds the output bias into `decoder`, it has no separate bias parameter.+        del self.bias+++@auto_docstring+class GteForMaskedLM(JinaEmbeddingsV3ForMaskedLM):+    _tied_weights_keys = {"lm_head.decoder.weight": "gte.embeddings.word_embeddings.weight"}++    def __init__(self, config: GteConfig):+        PreTrainedModel.__init__(self, config)++        self.gte = GteModel(config, add_pooling_layer=False)+        self.lm_head = GteLMHead(config)++        # Initialize weights and apply final processing+        self.post_init()++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.FloatTensor | None = None,+        token_type_ids: torch.LongTensor | None = None,+        position_ids: torch.LongTensor | None = None,+        inputs_embeds: torch.FloatTensor | None = None,+        labels: torch.LongTensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | MaskedLMOutput:+        # Unlike JinaEmbeddingsV3ForMaskedLM, the backbone is bound as `self.gte`, which modular cannot rename.+        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]`+        """+        outputs = self.gte(+            input_ids=input_ids,+            attention_mask=attention_mask,+            token_type_ids=token_type_ids,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            return_dict=True,+            **kwargs,+        )+        sequence_output = outputs[0]+        prediction_scores = self.lm_head(sequence_output)++        masked_lm_loss = None+        if labels is not None:+            # move labels to correct device+            labels = labels.to(prediction_scores.device)+            loss_fct = CrossEntropyLoss()  # -100 index = padding token+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))++        return MaskedLMOutput(+            loss=masked_lm_loss,+            logits=prediction_scores,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++class GteForSequenceClassification(BertForSequenceClassification):+    def __init__(self, config: GteConfig):+        super().__init__(config)+        self.gte = GteModel(config, add_pooling_layer=True)+++class GteForTokenClassification(BertForTokenClassification):+    pass+++__all__ = [+    "GteConfig",+    "GtePreTrainedModel",+    "GteModel",+    "GteForMaskedLM",+    "GteForSequenceClassification",+    "GteForTokenClassification",+]

GitHub publishes no diff for this file. View it on GitHub ↗

@@ -0,0 +1,385 @@+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++import unittest++from transformers import (+    AutoModel,+    AutoModelForMaskedLM,+    AutoModelForSequenceClassification,+    AutoTokenizer,+    GteConfig,+    is_torch_available,+)+from transformers.testing_utils import Expectations, require_torch, slow, torch_device++from ...test_configuration_common import ConfigTester+from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask+from ...test_pipeline_mixin import PipelineTesterMixin+++if is_torch_available():+    import torch++    from transformers import (+        GteForMaskedLM,+        GteForSequenceClassification,+        GteForTokenClassification,+        GteModel,+    )+++class GteModelTester:+    def __init__(+        self,+        parent,+        batch_size=13,+        seq_length=7,+        is_training=True,+        use_input_mask=True,+        use_token_type_ids=True,+        use_labels=True,+        vocab_size=99,+        hidden_size=32,+        num_hidden_layers=2,+        num_attention_heads=4,+        intermediate_size=37,+        hidden_act="gelu",+        hidden_dropout_prob=0.1,+        attention_probs_dropout_prob=0.1,+        max_position_embeddings=512,+        type_vocab_size=16,+        type_sequence_label_size=2,+        initializer_range=0.02,+        num_labels=3,+        num_choices=4,+        scope=None,+    ):+        self.parent = parent+        self.batch_size = batch_size+        self.seq_length = seq_length+        self.is_training = is_training+        self.use_input_mask = use_input_mask+        self.use_token_type_ids = use_token_type_ids+        self.use_labels = use_labels+        self.vocab_size = vocab_size+        self.hidden_size = hidden_size+        self.num_hidden_layers = num_hidden_layers+        self.num_attention_heads = num_attention_heads+        self.intermediate_size = intermediate_size+        self.hidden_act = hidden_act+        self.hidden_dropout_prob = hidden_dropout_prob+        self.attention_probs_dropout_prob = attention_probs_dropout_prob+        self.max_position_embeddings = max_position_embeddings+        self.type_vocab_size = type_vocab_size+        self.type_sequence_label_size = type_sequence_label_size+        self.initializer_range = initializer_range+        self.num_labels = num_labels+        self.num_choices = num_choices+        self.scope = scope++    def prepare_config_and_inputs(self):+        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)++        input_mask = None+        if self.use_input_mask:+            input_mask = random_attention_mask([self.batch_size, self.seq_length])++        token_type_ids = None+        if self.use_token_type_ids:+            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)++        sequence_labels = None+        token_labels = None+        choice_labels = None+        if self.use_labels:+            sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)+            token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)+            choice_labels = ids_tensor([self.batch_size], self.num_choices)++        config = self.get_config()++        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels++    def get_config(self):+        """+        Returns a tiny configuration by default.+        """+        return GteConfig(+            vocab_size=self.vocab_size,+            hidden_size=self.hidden_size,+            num_hidden_layers=self.num_hidden_layers,+            num_attention_heads=self.num_attention_heads,+            intermediate_size=self.intermediate_size,+            hidden_act=self.hidden_act,+            hidden_dropout_prob=self.hidden_dropout_prob,+            attention_probs_dropout_prob=self.attention_probs_dropout_prob,+            max_position_embeddings=self.max_position_embeddings,+            type_vocab_size=self.type_vocab_size,+            initializer_range=self.initializer_range,+        )++    def create_and_check_model(+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels+    ):+        model = GteModel(config=config)+        model.to(torch_device)+        model.eval()+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)+        result = model(input_ids, token_type_ids=token_type_ids)+        result = model(input_ids)+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))+        self.parent.assertEqual(result.pooler_output, None)++    def create_and_check_model_without_token_types(+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels+    ):+        config.type_vocab_size = 0+        model = GteModel(config=config)+        model.to(torch_device)+        model.eval()+        self.parent.assertIsNone(model.embeddings.token_type_embeddings)+        result = model(input_ids, attention_mask=input_mask)+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))++    def create_and_check_for_masked_lm(+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels+    ):+        model = GteForMaskedLM(config=config)+        model.to(torch_device)+        model.eval()+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))++    def create_and_check_for_sequence_classification(+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels+    ):+        config.num_labels = self.num_labels+        model = GteForSequenceClassification(config)+        model.to(torch_device)+        model.eval()+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))++    def create_and_check_for_token_classification(+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels+    ):+        config.num_labels = self.num_labels+        model = GteForTokenClassification(config=config)+        model.to(torch_device)+        model.eval()+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))++    def prepare_config_and_inputs_for_common(self):+        config, input_ids, token_type_ids, input_mask, _, _, _ = self.prepare_config_and_inputs()+        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}+        return config, inputs_dict+++@require_torch+class GteModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):+    all_model_classes = (+        (+            GteModel,+            GteForMaskedLM,+            GteForSequenceClassification,+            GteForTokenClassification,+        )+        if is_torch_available()+        else ()+    )+    pipeline_model_mapping = (+        {+            "feature-extraction": GteModel,+            "fill-mask": GteForMaskedLM,+            "text-classification": GteForSequenceClassification,+            "token-classification": GteForTokenClassification,+            "zero-shot": GteForSequenceClassification,+        }+        if is_torch_available()+        else {}+    )+    model_split_percents = [0.5, 0.8, 0.9]++    def setUp(self):+        self.model_tester = GteModelTester(self)+        self.config_tester = ConfigTester(self, config_class=GteConfig, hidden_size=37)++    def test_config(self):+        self.config_tester.run_common_tests()++    def test_model(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_model(*config_and_inputs)++    def test_model_without_token_types(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_model_without_token_types(*config_and_inputs)++    def test_for_masked_lm(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)++    def test_for_sequence_classification(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_for_sequence_classification(*config_and_inputs)++    def test_for_token_classification(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_for_token_classification(*config_and_inputs)+++@require_torch+class GteModelIntegrationTest(unittest.TestCase):+    sentences = ["Plants create oxygen.", "Photosynthesis is a process where plants create oxygen."]++    @slow+    def test_inference_no_head_multilingual(self):+        model = AutoModel.from_pretrained("alibaba-nlp-community/gte-multilingual-base", dtype=torch.float32).to(+            torch_device+        )+        tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")++        inputs = tokenizer(self.sentences, return_tensors="pt", padding=True, truncation=True).to(torch_device)++        with torch.no_grad():+            output = model(**inputs)[0]++        expected_shape = torch.Size((2, 15, 768))+        self.assertEqual(output.shape, expected_shape)++        # fmt: off+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [[ 0.2660, -0.2507,  0.3888], [ 0.5936, -0.4777,  0.5725], [ 0.5867, -0.4302,  0.5467]],+                        [[ 0.4506, -0.2547,  0.4193], [ 0.3973, -0.2549,  0.3836], [ 0.3309, -0.3932,  0.4126]],+                    ]+                ),+            }+        ).get_expectation()+        # fmt: on++        torch.testing.assert_close(output[:, 1:4, 1:4].cpu().detach(), expected_slice, rtol=1e-3, atol=1e-3)++    @slow+    def test_inference_no_head_english_v1_5(self):+        model = AutoModel.from_pretrained("alibaba-nlp-community/gte-base-en-v1.5", dtype=torch.float32).to(+            torch_device+        )+        tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-base-en-v1.5")++        inputs = tokenizer(self.sentences, return_tensors="pt", padding=True, truncation=True).to(torch_device)++        with torch.no_grad():+            output = model(**inputs)[0]++        expected_shape = torch.Size((2, 13, 768))+        self.assertEqual(output.shape, expected_shape)++        # fmt: off+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [[-0.0137, -0.4422,  0.5301], [-0.2834,  0.4779,  0.6164], [-0.2570,  0.3076,  0.7647]],+                        [[-0.1446, -0.0088, -0.1134], [-0.1478,  0.1888,  0.3145], [-0.1605,  0.1750,  0.1975]],+                    ]+                ),+            }+        ).get_expectation()+        # fmt: on++        torch.testing.assert_close(output[:, 1:4, 1:4].cpu().detach(), expected_slice, rtol=1e-3, atol=1e-3)++    @slow+    def test_inference_no_head_snowflake_arctic_embed(self):+        model = AutoModel.from_pretrained("Snowflake/snowflake-arctic-embed-m-v2.0", dtype=torch.float32).to(+            torch_device+        )+        tokenizer = AutoTokenizer.from_pretrained("Snowflake/snowflake-arctic-embed-m-v2.0")++        inputs = tokenizer(self.sentences, return_tensors="pt", padding=True, truncation=True).to(torch_device)++        with torch.no_grad():+            output = model(**inputs)[0]++        expected_shape = torch.Size((2, 15, 768))+        self.assertEqual(output.shape, expected_shape)++        # fmt: off+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [[ 1.2827,  0.5775, -0.4335], [ 1.3125,  0.2771, -0.4684], [ 1.4864,  0.1338, -0.3023]],+                        [[ 1.6184, -0.2137, -0.8186], [ 1.6171, -0.2500, -0.5360], [ 1.4568, -0.3333, -0.5992]],+                    ]+                ),+            }+        ).get_expectation()+        # fmt: on++        torch.testing.assert_close(output[:, 1:4, 1:4].cpu().detach(), expected_slice, rtol=1e-3, atol=1e-3)++    @slow+    def test_inference_reranker(self):+        model = AutoModelForSequenceClassification.from_pretrained(+            "alibaba-nlp-community/gte-multilingual-reranker-base", dtype=torch.float32+        ).to(torch_device)+        tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-reranker-base")++        inputs = tokenizer(self.sentences, return_tensors="pt", padding=True, truncation=True).to(torch_device)++        with torch.no_grad():+            output = model(**inputs).logits++        expected_shape = torch.Size((2, 1))+        self.assertEqual(output.shape, expected_shape)++        expected_slice = Expectations({(None, None): torch.tensor([[0.2428], [0.5710]])}).get_expectation()++        torch.testing.assert_close(output.cpu().detach(), expected_slice, rtol=1e-3, atol=1e-3)++    @slow+    def test_inference_masked_lm(self):+        model = AutoModelForMaskedLM.from_pretrained(+            "alibaba-nlp-community/gte-multilingual-mlm-base", dtype=torch.float32+        ).to(torch_device)+        tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-mlm-base")++        inputs = tokenizer(self.sentences, return_tensors="pt", padding=True, truncation=True).to(torch_device)++        with torch.no_grad():+            output = model(**inputs).logits++        expected_shape = torch.Size((2, 15, 250048))+        self.assertEqual(output.shape, expected_shape)++        # fmt: off+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [[-1.9263,  5.8634,  3.9742], [-1.2569,  9.1684,  6.1088], [-1.1027,  6.3202,  8.6400]],+                        [[-2.3521,  4.1811,  9.3579], [-2.2531, -0.8155,  5.1775], [-2.1518,  4.4150,  7.3230]],+                    ]+                ),+            }+        ).get_expectation()+        # fmt: on++        torch.testing.assert_close(output[:, 1:4, 1:4].cpu().detach(), expected_slice, rtol=1e-3, atol=1e-3)