model: Add GTE to Transformers - #48416
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Model linter — first pass
transformers-mlinter found 1 item(s) in the model files this PR touches. These are structural conventions a maintainer would otherwise flag by hand.
This is automated and advisory — it does not block merging.
| rule | count | what it checks |
|---|---|---|
TRF041 |
1 | A config-gated branch must carry a # CODEPATH: note saying which checkpoints diverge. |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Looks very solid, biggest blocker are the conversion needed for the configs atm -> let's rather make PRs on the hub and use the revision to point to it
|
|
||
| *Driven by the simple and effective Dense O2O, DEIM demonstrates faster convergence and enhanced performance. In this work, we extend it with DINOv3 features, resulting in DEIMv2. DEIMv2 spans eight model sizes from X to Atto, covering GPU, edge, and mobile deployment. For the X, L, M, and S variants, we adopt DINOv3-pretrained / distilled backbones and introduce a Spatial Tuning Adapter (STA), which efficiently converts DINOv3's single-scale output into multi-scale features and complements strong semantics with fine-grained details to enhance detection. For ultra-lightweight models (Nano, Pico, Femto, and Atto), we employ HGNetv2 with depth and width pruning to meet strict resource budgets. Together with a simplified decoder and an upgraded Dense O2O, this unified design enables DEIMv2 to achieve a superior performance-cost trade-off across diverse scenarios, establishing new state-of-the-art results. Notably, our largest model, DEIMv2-X, achieves 57.8 AP with only 50.3M parameters, surpassing prior X-scale models that require over 60M parameters for just 56.5 AP. On the compact side, DEIMv2-S is the first sub-10M model (9.71M) to exceed the 50 AP milestone on COCO, reaching 50.9 AP. Even the ultra-lightweight DEIMv2-Pico, with just 1.5M parameters, delivers 38.5 AP-matching YOLOv10-Nano (2.3M) with ~50% fewer parameters.* | ||
|
|
||
| This model was contributed by [Harshal Janjani](https://huggingface.co/harshaljanjani). |
There was a problem hiding this comment.
let's move it to a different PR but definitely merging then to credit!
|
|
||
| ## Usage examples | ||
|
|
||
| Embeddings are taken from the `[CLS]` token and normalized. |
There was a problem hiding this comment.
Is this not for sequence classification then? Or am I missing something?
Might be nicer to just follow the other bert likes a bit with pipeline usage etc
There was a problem hiding this comment.
So [CLS] is being used for pooling into an embedding, not classification here, as the model was trained. Classification is a separate model path with an added head. Rewrote it like bert and jina_embeddings_v3
| Fine-tuning uses the standard forward and backward pass: | ||
|
|
||
| ```python | ||
| import torch | ||
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | ||
|
|
||
| model_id = "harshaljanjani/gte-multilingual-base-hf" | ||
| tokenizer = AutoTokenizer.from_pretrained(model_id) | ||
| model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2, device_map="auto") | ||
|
|
||
| inputs = tokenizer(["a positive review", "a negative review"], padding=True, return_tensors="pt").to(model.device) | ||
| labels = torch.tensor([1, 0], device=model.device) | ||
|
|
||
| loss = model(**inputs, labels=labels).loss | ||
| loss.backward() | ||
| ``` | ||
|
|
||
| The model is compatible with [`torch.compile`]: | ||
|
|
||
| ```python | ||
| import torch | ||
| from transformers import AutoModel, AutoTokenizer | ||
|
|
||
| model_id = "harshaljanjani/gte-multilingual-base-hf" | ||
| tokenizer = AutoTokenizer.from_pretrained(model_id) | ||
| model = AutoModel.from_pretrained(model_id, device_map="auto") | ||
| model = torch.compile(model) | ||
|
|
||
| inputs = tokenizer("what is the capital of China?", return_tensors="pt").to(model.device) | ||
| with torch.no_grad(): | ||
| outputs = model(**inputs) | ||
| ``` |
There was a problem hiding this comment.
I think those make more sense for e.g. audio models as they are more special but here we can focus on more casual usage with auto and pipeline
|
|
||
| - 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/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. |
There was a problem hiding this comment.
Will see but ig we register in post init properly?
There was a problem hiding this comment.
Here, with the Hub PRs the configs now carry rope_parameters directly as you suggested :)
| rope_theta = kwargs.pop("rope_theta", self.default_theta) | ||
|
|
||
| # GTE's static NTK scaling is exactly a linear scaling of `base * factor`. | ||
| if rope_scaling is not None and rope_scaling["type"] == "ntk": |
There was a problem hiding this comment.
Can we link an example for which model
There was a problem hiding this comment.
If we are already submitting PRs with rope parameters directly, then we can directly do that there instead of overriding this method, no?
| def _init_weights(self, module): | ||
| # None of the inherited buffer initialisations apply, GTE keeps no such buffers. | ||
| PreTrainedModel._init_weights(self, module) |
There was a problem hiding this comment.
lets use attribute error instead then to not inherit anything
There was a problem hiding this comment.
Tried it, the converter only pops class attributes so a method name raises KeyError: '_init_weights' here
| class GteForMaskedLM(JinaEmbeddingsV3ForMaskedLM): | ||
| _tied_weights_keys = {"lm_head.decoder.weight": "gte.embeddings.word_embeddings.weight"} | ||
|
|
||
| def __init__(self, config: GteConfig): |
There was a problem hiding this comment.
this override is not needed no? at max, maybe self.gte
There was a problem hiding this comment.
The parent assigns the backbone to self.roberta which the converter can't rename, so super().__init__ + self.gte creates both backbones :(
Same as nomic
| GteForMultipleChoice, | ||
| GteForQuestionAnswering, | ||
| GteForSequenceClassification, | ||
| GteForTokenClassification, |
There was a problem hiding this comment.
Imo we could reduce a bit as in nomic bert iirc. not all the for are super used so could make maintenance a bit easier
There was a problem hiding this comment.
Done, dropped ForQuestionAnswering and ForMultipleChoice so the set matches Nomic's four.
| class GteModelIntegrationTest(unittest.TestCase): | ||
| sentences = ["Plants create oxygen.", "Photosynthesis is a process where plants create oxygen."] | ||
|
|
||
| # TODO: Point these back at Alibaba-NLP and Snowflake once their configs declare `model_type: "gte"`. |
There was a problem hiding this comment.
If we have to submit conversions either way, we could also just add the rope parameters directly to the configs no?
There was a problem hiding this comment.
| sentences = ["Plants create oxygen.", "Photosynthesis is a process where plants create oxygen."] | ||
|
|
||
| # TODO: Point these back at Alibaba-NLP and Snowflake once their configs declare `model_type: "gte"`. | ||
| # NOTE: The upstream repos carry an `auto_map`, so `Auto*` resolves them to that remote code |
There was a problem hiding this comment.
Hmm but that shouldnt be the case 🤔 shouldnt we need to pass trust_remote_code=True for that to happen
You could use a revision to make both live at the same repo, i.e. add what we need for the transformers integration
There was a problem hiding this comment.
Ah, my note was misleading apologies! Also applied the revision param here and everywhere with TODOs
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Ok my only bigger ask is to just make the ntk remotely into rope parameters directly as well, the rest is smaller in nature
Afterwards, I will ask internally for contacts so we may merge the remote hub PRs 🫡
| task="feature-extraction", | ||
| model="Alibaba-NLP/gte-multilingual-base", | ||
| revision="refs/pr/31", | ||
| device=0 |
There was a problem hiding this comment.
| device=0 |
I think we use device map auto on pipelines so can be removed
| "Alibaba-NLP/gte-multilingual-base", | ||
| revision="refs/pr/31", | ||
| device_map="auto", | ||
| attn_implementation="sdpa" |
There was a problem hiding this comment.
| attn_implementation="sdpa" |
same as its default
|
|
||
| [[autodoc]] GteForSequenceClassification | ||
| - forward | ||
|
|
| rope_theta = kwargs.pop("rope_theta", self.default_theta) | ||
|
|
||
| # GTE's static NTK scaling is exactly a linear scaling of `base * factor`. | ||
| if rope_scaling is not None and rope_scaling["type"] == "ntk": |
There was a problem hiding this comment.
If we are already submitting PRs with rope parameters directly, then we can directly do that there instead of overriding this method, no?
| @torch.no_grad() | ||
| def _init_weights(self, module): | ||
| # None of the inherited buffer initialisations apply, GTE keeps no such buffers. | ||
| PreTrainedModel._init_weights(self, module) |
There was a problem hiding this comment.
| @torch.no_grad() | |
| def _init_weights(self, module): | |
| # None of the inherited buffer initialisations apply, GTE keeps no such buffers. | |
| PreTrainedModel._init_weights(self, module) | |
| def _init_weights(self, **super_kwargs): | |
| raise AttributeError("Uses base super call") |
pretty sure this should work, can you retry?
|
@harshaljanjani thanks a lot, just checked most PRs on the hub. Imo we could maybe simplify by quite a bit when we use That would handle most of the PRs need for writing rope parameters explicitly (and would reduce my attempts at contacting orgs 😬) Other than that, if we do not reach them we could still get them in even without the hub PRs
|
|
@vasqu A couple of clarifying questions. → I tried Please do let me know what you think about this! |
|
This is not to exchange every hub PR tbh but at least make it easier for us on a few models, e.g. as you mentioned NTK and different rope thetas likely still need their hub PRs For example, with this default, we can theoretically still work with the snowflake implementation without the hub PR |
|
I understand, thank you for the clarification, pushed the changes! The PRs are down from the previous 8 to just the 4 NTK ones. The PRs that are important now: → https://huggingface.co/Alibaba-NLP/gte-multilingual-base/discussions/31 |
| code_revision = kwargs.pop("code_revision", None) | ||
|
|
||
| config_dict, unused_kwargs = PreTrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs) | ||
| # GTE checkpoints ship the `new` model type, and the ones without NTK scaling need nothing else to load natively |
There was a problem hiding this comment.
Yea no please revert this, we definitely need to update those models at least on the hub
There was a problem hiding this comment.
Reverted, so the PRs that become important now are:
→ https://huggingface.co/Alibaba-NLP/gte-multilingual-base/discussions/31
→ https://huggingface.co/Alibaba-NLP/gte-base-en-v1.5/discussions/17
→ https://huggingface.co/Alibaba-NLP/gte-large-en-v1.5/discussions/26
→ https://huggingface.co/Alibaba-NLP/gte-multilingual-mlm-base/discussions/2
→ https://huggingface.co/Alibaba-NLP/gte-multilingual-reranker-base/discussions/23
→ https://huggingface.co/Alibaba-NLP/gte-en-mlm-base/discussions/3
→ https://huggingface.co/Alibaba-NLP/gte-en-mlm-large/discussions/4
Basically the Alibaba-NLP ones not Snowflake. Snowflake as you pointed out runs on main :)
There was a problem hiding this comment.
Perf, trying to contact the alibaba team so gotta be a bit patient for now 🤗
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, gte |
CI recapDashboard: View test results in Grafana |
2 participants
What does this PR do?
→ This PR adds GTE to Transformers!
→ Completes the GTE and Snowflake GTE items in #42738 (
Snowflake/snowflake-arctic-embed-m-v2.0ships the same architecture)→ Model Checkpoints
→ Original Implementation
→ Paper
cc: @vasqu
Code Agent Policy
Before submitting
🤖 mlinter review state
fdd2305dcf6163f0