Dewey En Beta

🚀 Dewey長上下文嵌入模型：技術報告

本技術報告介紹了Dewey，這是一種新穎的長上下文嵌入模型，旨在提升長文檔場景下的檢索性能。該模型基於ModernBERT架構，結合基於指令的訓練方法，使嵌入與特定任務需求相匹配。Dewey具有128k的上下文窗口、多向量表示和靈活的分塊機制，在LongEmbed基準測試中取得了優異的成績。

🚀 快速開始

本模型由Richinfo合作發佈，採用了一種新穎的訓練方法。雖然尚未完全理解其底層原理，但已取得了不錯的效果。因此，我們決定將該模型開源，並希望有人能對其進行測試並提供反饋！

技術報告鏈接：https://arxiv.org/abs/2503.20376

該模型的核心訓練方法將在NovaSearch團隊開源的RAG-Retrieval倉庫中實現，歡迎star！

本模型基於answerdotai/ModernBERT-large，感謝他們的分享！

✨ 主要特性

1. 長文本支持：最大長度為128k，參數大小為395M，僅支持英文。
2. 多向量表示：支持單向量和多向量（類似於Colbert，但向量數量更少，僅為標記數量的0.5%）。
3. 短文本性能優異：在短文本評估（MTEB-eng-v2）中取得了令人印象深刻的結果，未使用MTEB訓練集，甚至超越了幾個7B大小的模型。
4. 長文本表現出色：在長文本評估LongEmbed中，單向量超越了許多大型和商業模型。若使用多向量，平均得分排名第一。目前，我們的得分為0.86，而當前第一名的得分為0.79。
5. 編碼速度快：得益於ModernBert的架構優勢，長文本的編碼速度仍然非常快。
6. 靈活的多向量組合：多向量可以理解為跨度或分塊級別，而非標記級別，因此可以根據自己的場景完全自定義分塊方式，非常靈活。

📦 安裝指南

本項目未提供具體安裝命令，可參考以下步驟：

1. 克隆項目倉庫：

git clone https://github.com/NovaSearch-Team/RAG-Retrieval.git
cd RAG-Retrieval

2. 安裝依賴：

pip install -r requirements.txt

💻 使用示例

基礎用法

我們建議結合模型架構圖閱讀以下內容。

請仔細閱讀modeling_dewey_v1.py和custom_st.py，這些代碼易於閱讀，會對你有很大幫助！

# 單向量使用示例
import os

# os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import torch
from sentence_transformers import SentenceTransformer

RETRIEVE_Q_PROMPT = "<|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>"
RETRIEVE_P_PROMPT = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"
model = SentenceTransformer(
    "infgrad/dewey_en_beta",
    trust_remote_code=True,
    model_kwargs={
        "torch_dtype": torch.bfloat16,
        "attn_implementation": "flash_attention_2"
    },
    config_kwargs={"single_vector_type": "mean"}
).cuda().bfloat16().eval()
# the choice of single_vector_type:
## for short text (<1k): cls_add_mean
## for long text (>1k): mean

# the max length of model is 128*1024
model.max_seq_length = 32 * 1024

query_vectors = model.encode(
    sentences=[f"{RETRIEVE_Q_PROMPT}What is a computer composed of?", f"{RETRIEVE_Q_PROMPT}why the sky is blue"]
)
passage_vectors = model.encode(
    sentences=[
        f"{RETRIEVE_P_PROMPT}Central processing unit (CPU), memory (RAM), storage (hard drive or SSD), input/output devices (keyboard, mouse, monitor), and a motherboard",
        f"{RETRIEVE_P_PROMPT}Shorter wavelengths of light, such as blue and violet, are scattered more by gases and particles in Earth's atmosphere.",
    ]
)

print(query_vectors @ passage_vectors.T)
# the output is:
# [[0.52512825 0.19771025]
#  [0.17617573 0.5918883 ]]

高級用法

自動分塊獲取多向量

import os
import numpy as np

# os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from pydantic import BaseModel
from typing import Optional, List
from transformers import AutoTokenizer, AutoModel


class TextSpan(BaseModel):
    s: int
    e: int
    text: Optional[str] = None
    module_name: str


RETRIEVE_Q_PROMPT = "<|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>"
RETRIEVE_P_PROMPT = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"
model = AutoModel.from_pretrained(
    "infgrad/dewey_en_beta",
    trust_remote_code=True,
    attn_implementation="flash_attention_2"
).cuda().bfloat16()
model.tokenizer = AutoTokenizer.from_pretrained("infgrad/dewey_en_beta")
max_seq_length = 32 * 1024

q_list = ["why the sky is blue"]
p_list = [
    """
    I’ve been trying to understand why the sky changes colors, and I think I understand most of it, but something in the online explanations doesn’t make it clear for me:

I’ve read:

sky is blue because blue light gets scattered the most during the day.

in the evening it turns red because now even more of the blue light gets scattered

So a few questions:

The scattering of light during the day: does it mean that blue light gets reflected off air particles and reaches our eyes, while the rest of the frequencies pass through and reach the ground?

Surely some of the other frequencies also get scattered during the day, just in much smaller amounts?

So during the evening blue light gets scattered even more, to the point where even less of it reaches the eyes?

And so it gets red because now we can see the lower frequencies being scattered without blue overshadowing them?\

Trying to word it myself: during the day only the highest frequencies get filtered, but during the evening also lower frequencies get filtered, because now the “light strainer” (air) is just catching more of it?\

It gets darker in the evening without a good ability to see colors because there’s is no blue and so on light to reflect off of objects?\

Is it ok to speak about light as a frequency? Or it’s only correct to say “wave length”?

Blue light is scattered in all directions by the tiny molecules of air in Earth's atmosphere. Blue is scattered more than other colors because it travels as shorter, smaller waves. 
This is why we see a blue sky most of the time. Closer to the horizon, the sky fades to a lighter blue or white.
    """
]

# query should be a single vector, so we set chunk_size as -1 to avoid chunk.
# If chunk size is -1, the model will return an array with shape of (2,2048) consisting of cls_vector and mean_vector(mean of all token embeddings).
query_vectors = model.encode(
    sentences=q_list,
    use_cuda=True,
    show_progress_bar=True,
    chunk_size=-1,
    chunk_overlap=32,
    convert_to_tensor=False,
    max_seq_length=max_seq_length,
    batch_size=8,
    normalize_embeddings=True,
    prompt=RETRIEVE_Q_PROMPT,
    fast_chunk=False

)[0]
# query vector do not need multi vector, we only use mean as final single vector
pred = [vecs[1:2, :] for vecs in query_vectors]

# spans_list contail each chunk's span, you can use span to get text
spans_list: List[List[TextSpan]]
passage_vectors_list: List[np.ndarray]
passage_vectors_list, spans_list = model.encode(
    sentences=p_list,
    use_cuda=True,
    show_progress_bar=True,
    chunk_size=64,
    chunk_overlap=8,
    convert_to_tensor=False,
    max_seq_length=max_seq_length,
    batch_size=8,
    normalize_embeddings=True,
    prompt=RETRIEVE_P_PROMPT,
    fast_chunk=True,  # if fast_chunk is true, directly chunk on input ids, else using RecursiveCharacterTextSplitter
)
# spans_list stores each passage's spans, passage_vectors_list stores each passage's vectors so len(spans_list) == len(p_list) and len(spans_list) == len(passage_vectors_list)
# for a passage's spans and vectors, each span corresponds to a vector (1*2048). So, len(spans_list[idx]) ==  len(passage_vectors_list[idx])
print((query_vectors[0] @ passage_vectors_list[0].T).max())
# output 0.7331543
# get each chunk's content
for spans, passage in zip(spans_list, p_list):
    text_ids = model.tokenizer.encode(RETRIEVE_P_PROMPT + passage)
    for span in spans:
        s, e = span.s, span.e
        chunk_text = model.tokenizer.decode(
            text_ids[s:e],
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True
        ).strip()

手動分塊獲取多向量

import os
import numpy as np

# os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from pydantic import BaseModel
from typing import Optional, List
from transformers import AutoTokenizer, AutoModel


class TextSpan(BaseModel):
    s: int
    e: int
    text: Optional[str] = None
    module_name: str


prompt = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"

# load model
model = AutoModel.from_pretrained(
    "infgrad/dewey_en_beta",
    trust_remote_code=True,
    attn_implementation="flash_attention_2"
)
model.tokenizer = AutoTokenizer.from_pretrained("infgrad/dewey_en_beta")
max_seq_length = 32 * 1024

# chunk text
passage = "this sentence 1. this sentence 2. this sentence 3"
chunks = ["this sentence 1. this sentence 2.", "this sentence 2. this sentence 3"]
prompt_length = len(model.tokenizer.tokenize(prompt))
text_spans = [
    # s=0, e=1 means that this vector is cls vector, so the module_name is cls_linear, otherwise the module_name is chunk_linear
    TextSpan(s=0, e=1, module_name="cls_linear")
]
for chunk in chunks:
    s = passage.find(chunk)
    e = s + len(chunk)
    text_spans.append(
        TextSpan(
            # add 1, as there is a [CLS] token at the beginning of text.
            s=1 + prompt_length + len(model.tokenizer.tokenize(passage[:s])),
            e=1 + prompt_length + len(model.tokenizer.tokenize(passage[:e])),
            module_name="chunk_linear"
        )
    )

spans_list: List[List[TextSpan]]
passage_vectors_list: List[np.ndarray]
passage_vectors_list, _ = model.encode(
    sentences=[passage],
    use_cuda=False,
    show_progress_bar=True,
    chunk_size=64,
    chunk_overlap=12,
    convert_to_tensor=False,
    max_seq_length=max_seq_length,
    batch_size=8,
    normalize_embeddings=True,
    prompt=prompt,
    fast_chunk=True,
    batch_text_spans=[text_spans]
)
print(passage_vectors_list[0].shape, passage_vectors_list[0][:, 2])
# the output is (3, 2048) [0.01461297 0.02085092 0.0022509 ]

📚 詳細文檔

MTEB(eng, v2)

• 評估鏈接：http://mteb-leaderboard.hf.space/?benchmark_name=MTEB%28eng%2C+v2%29
• 復現腳本：https://huggingface.co/infgrad/dewey_en_beta/blob/main/scripts/evaluate/run_evaluate_mteb_dewey_en_beta.py

LongEmbed

• 評估鏈接：http://mteb-leaderboard.hf.space/?benchmark_name=LongEmbed
• 復現腳本：https://huggingface.co/infgrad/dewey_en_beta/blob/main/scripts/evaluate/run_evaluate_long_embed.py

LoCoV1

• 評估鏈接：
  • https://huggingface.co/datasets/hazyresearch/LoCoV1-Queries
  • https://huggingface.co/datasets/hazyresearch/LoCoV1-Documents
• 復現腳本：https://huggingface.co/infgrad/dewey_en_beta/blob/main/scripts/evaluate/run_evaluate_loco.py
• 評估指標：NDCG@10

🔧 技術細節

本模型基於answerdotai/ModernBERT-large，採用了一種新穎的訓練方法。該模型具有以下技術特點：

1. 長上下文處理：支持最大長度為128k的文本，能夠處理長文檔。
2. 多向量表示：支持單向量和多向量表示，多向量表示可以提高檢索的準確性。
3. 靈活的分塊機制：可以根據需要自定義分塊方式，以適應不同的應用場景。
4. 快速編碼：得益於ModernBert的架構優勢，模型的編碼速度非常快。

📄 許可證

本模型採用MIT許可證。

🔧 侷限性

• 語言限制：僅支持英文文本。
• 短文本性能：在短文本任務上，性能可能不如傳統的短文本嵌入模型。
• 模型階段：該模型仍處於alpha或beta階段，可能會有一些意外的行為。

📖 引用

@misc{zhang2025deweylongcontextembedding,
      title={Dewey Long Context Embedding Model: A Technical Report}, 
      author={Dun Zhang and Panxiang Zou and Yudong Zhou},
      year={2025},
      eprint={2503.20376},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2503.20376}, 
}