stella_en

Modifications for transformers >= 5

transformers 5.x instantiates models on the meta device, which broke this model in two independent, silent ways (no errors — just garbage embeddings):

1. Rotary caches scrambled — the non-persistent inv_freq/cos/sin buffers computed in __init__ are materialized from the meta device as uninitialized memory and were never recomputed, corrupting every position. Fixed in modeling_qwen.py: the rotary caches are now rebuilt lazily on the first real forward pass. No behavior change on transformers 4.x.
2. Checkpoint weights never load — the 5.x loader reports a clean load (0 missing keys) for this old-style remote-code class but leaves every parameter randomly initialized. This cannot be fixed from the modeling file. Until it is fixed upstream in transformers, reload the state dict manually after loading:

import torch
from sentence_transformers import SentenceTransformer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

model = SentenceTransformer("it-just-works/stella_en_1.5B_v5_bf16", trust_remote_code=True,
                            model_kwargs={"dtype": torch.bfloat16})
sd = load_file(hf_hub_download("it-just-works/stella_en_1.5B_v5_bf16", "model.safetensors"))
sd = {k.removeprefix("model."): v.to(torch.bfloat16) for k, v in sd.items()}
model[0].auto_model.load_state_dict(sd, strict=False)

Sanity check after loading: encode two related and one unrelated sentence — the related pair must clearly out-score the unrelated one. If everything lands at cosine ~0 (or ~0.95 for all pairs), the model is silently broken.

Modifications for transformers 4.45+

• Cache.get_usable_length() was removed from the transformers cache API; modeling_qwen.py now ships a _get_usable_past_kv_length() helper used everywhere the old method was called, restoring cache compatibility.
• Weights converted to bfloat16 (convert_to_bf16.py included); ONNX variants removed.

Updates

We released a Jasper and Stella model technology report and code.(2025.1)

Report: https://arxiv.org/abs/2412.19048

Codes: https://github.com/NLPJCL/RAG-Retrieval

Introduction

The models are trained based on Alibaba-NLP/gte-large-en-v1.5 and Alibaba-NLP/gte-Qwen2-1.5B-instruct. Thanks for their contributions!

We simplify usage of prompts, providing two prompts for most general tasks, one is for s2p, another one is for s2s.

Prompt of s2p task(e.g. retrieve task):

Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: {query}

Prompt of s2s task(e.g. semantic textual similarity task):

Instruct: Retrieve semantically similar text.\nQuery: {query}

The models are finally trained by MRL, so they have multiple dimensions: 512, 768, 1024, 2048, 4096, 6144 and 8192.

The higher the dimension, the better the performance. Generally speaking, 1024d is good enough. The MTEB score of 1024d is only 0.001 lower than 8192d.

Model directory structure

The model directory structure is very simple, it is a standard SentenceTransformer directory with a series of 2_Dense_{dims} folders, where dims represents the final vector dimension.

For example, the 2_Dense_256 folder stores Linear weights that convert vector dimensions to 256 dimensions. Please refer to the following chapters for specific instructions on how to use them.

Usage

You can use SentenceTransformers or transformers library to encode text.

Sentence Transformers

from sentence_transformers import SentenceTransformer

# This model supports two prompts: "s2p_query" and "s2s_query" for sentence-to-passage and sentence-to-sentence tasks, respectively.
# They are defined in `config_sentence_transformers.json`
query_prompt_name = "s2p_query"
queries = [
    "What are some ways to reduce stress?",
    "What are the benefits of drinking green tea?",
]
# docs do not need any prompts
docs = [
    "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.",
    "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.",
]

# ！The default dimension is 1024, if you need other dimensions, please clone the model and modify `modules.json` to replace `2_Dense_1024` with another dimension, e.g. `2_Dense_256` or `2_Dense_8192` !
model = SentenceTransformer("dunzhang/stella_en_1.5B_v5", trust_remote_code=True).cuda()
query_embeddings = model.encode(queries, prompt_name=query_prompt_name)
doc_embeddings = model.encode(docs)
print(query_embeddings.shape, doc_embeddings.shape)
# (2, 1024) (2, 1024)

similarities = model.similarity(query_embeddings, doc_embeddings)
print(similarities)
# tensor([[0.8179, 0.2958],
#         [0.3194, 0.7854]])

Transformers

import os
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

query_prompt = "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "
queries = [
    "Wha