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}, 
}