base_model:

• answerdotai/ModernBERT-large datasets:
• BAAI/Infinity-Instruct
• HuggingFaceFW/fineweb-edu language:
• en license: mit pipeline_tag: feature-extraction tags:
• sentence-transformers
• transformers library_name: sentence-transformers

Dewey長文コンテキスト埋め込みモデル: 技術レポート

このモデルは論文で発表されました。

論文要約

論文の要約は以下の通りです：

本技術レポートでは、長文ドキュメントシナリオにおける検索性能を向上させるために設計された新しい長文コンテキスト埋め込みモデルDeweyを紹介します。Deweyは、長いシーケンスを効率的に処理できることで知られるModernBERTアーキテクチャを基盤とし、埋め込みを特定のタスク要件に合わせるための命令ベースのトレーニングアプローチを組み込んでいます。Deweyの主な特徴には、128kのコンテキストウィンドウ、粒度を改善するためのマルチベクトル表現、カスタマイズ可能なベクトル組み合わせを可能にする柔軟なチャンキングメカニズムが含まれます。DeweyをLongEmbedベンチマークで評価したところ、いくつかの大規模モデルを上回る最先端の結果を達成しました。さらに、様々なアプリケーションでのDeweyの採用と適応を容易にするために、包括的な使用例と実装詳細を提供します。

1 はじめに

Richinfoと協力して、このリリースモデルは新しいアプローチでトレーニングされました。私たちはまだ根本的な原理を完全に理解していませんが、有望な結果を得たため、モデルをオープンソース化し、誰かがモデルをテストしてフィードバックを提供してくれることを願っています！

技術レポート: https://arxiv.org/abs/2503.20376

このモデルのコアトレーニング方法は、NovaSearchチームによってオープンソース化されたRAG-Retrievalリポジトリで実装される予定です。ぜひスターをお願いします！

このモデルはanswerdotai/ModernBERT-largeを基にしています。素晴らしいモデルを共有してくれたことに感謝します！

この埋め込みモデルには以下の特徴があります：

1. 最大長は128k、パラメータサイズは395Mで、英語のみをサポートします。
2. シングルベクトルとマルチベクトルの両方をサポート（Colbertに似ていますが、ベクトル数は少なく、トークン数の0.5%のみ）。
3. 短いテキスト評価（MTEB-eng-v2）で非常に印象的な結果を達成し、MTEBトレーニングセットを使用せず、いくつかの7Bサイズのモデルを上回りました。
4. 長文評価LongEmbedでは、シングルベクトルが多くの大規模および商用モデルを上回ります。マルチベクトルを使用すると、平均スコアが1位になります。現在、私たちのスコアは0.86で、現在の1位のスコアは0.79です。
5. 超高速なエンコーディング速度。ModernBertのアーキテクチャの利点により、長文のエンコーディング速度も非常に高速です。
6. 超柔軟なマルチベクトル組み合わせ方法。マルチベクトルはspanまたはchunkレベルと理解でき、トークンレベルではないため、チャンクの指定方法は完全にシナリオに応じてカスタマイズ可能で、非常に柔軟です。

2 使用方法

以下の内容はモデルアーキテクチャ図とともに読むことをお勧めします。

modeling_dewey_v1.pyとcustom_st.pyを注意深く読むことを強くお勧めします。これらのコードは読みやすく、多くの助けになります！

2.1 プロンプト

私たちのモデルは一種の命令埋め込みモデルであり、モデルを使用する際にはテキストの前にプロンプトを追加する必要があります。

検索タスクの場合、提供されたプロンプトを必ず使用してください：
クエリ: <|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>
パッセージ: <|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>

STSタスクの場合、提供されたプロンプトを必ず使用してください：
<|START_INSTRUCTION|>Generate semantically similar text<|END_INSTRUCTION|>

分類およびクラスタリングタスクの場合、独自のプロンプトを設計する必要があります。以下はいくつかの例です：
<|START_INSTRUCTION|>Classify text into intents<|END_INSTRUCTION|>
<|START_INSTRUCTION|>Classify text into toxic or not toxic<|END_INSTRUCTION|>
<|START_INSTRUCTION|>Output main category of Medrxiv papers based on the titles<|END_INSTRUCTION|>
<|START_INSTRUCTION|>Output topic or theme of news articles<|END_INSTRUCTION|>

2.2 シングルベクトル

シングルベクトルを使用する場合、私たちのモデルはSentenceTransformerと互換性があります。

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()
# single_vector_typeの選択：
## 短いテキスト（<1k）の場合: cls_add_mean
## 長いテキスト（>1k）の場合: mean

# モデルの最大長は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)
# 出力は：
# [[0.52512825 0.19771025]
#  [0.17617573 0.5918883 ]]

2.3 マルチベクトル

私たちのマルチベクトルはテキストスパン（つまりチャンク）に基づいているため、各ベクトルはコンテキストチャンクベクトルと見なすことができます。ドキュメントのマルチベクトルを取得するには、まずチャンクとそのスパンを取得する必要があります。

以下はマルチベクトルを取得する詳細な手順です：

ステップ1: ドキュメントをチャンクしてチャンクとスパンを取得します。これは私たちのencode関数を使用して行うことも、シナリオに応じて自分でドキュメントをチャンクすることもできます。
注意：自分でチャンクする場合、チャンクとスパンにプロンプトを含めないでください！
ステップ2: テキストをエンコードしてトークン埋め込みを取得します
ステップ3: スパン（つまり開始位置と終了位置）に基づいてチャンクベクトルを取得します。スパントークン埋め込みの平均をチャンクベクトルとして使用します（つまり、normalize(token_embed[start_position:end_position].mean(axis=0))）
ステップ4: 各スパンに対してステップ3を実行し、すべてのチャンクベクトルを取得します。また、span(0,1)とspan(1+prompt_len, text_len-1)を追加してグローバルベクトルを取得することもできます

検索タスクの場合、クエリベクトルはシングルベクトルである必要があるため、クエリとドキュメント間の最終スコアは、クエリと各ドキュメントベクトル間の最大スコアになります。
これはFAISS、MILVUSなどと互換性があります。トップkを拡大し、検索されたドキュメントに対して重複排除を行うだけです。

以下は詳細なコード例です。

2.3.1 encode関数内でテキストをチャンクする

マルチベクトルを取得するために、直接モデルのencodeメソッドを使用できます。
このメソッドは自動的にテキストをチャンクします。
fast_chunkパラメータを設定することでチャンク戦略を選択できます。fast_chunkがtrueの場合、入力ID上で直接チャンクし、それ以外の場合はRecursiveCharacterTextSplitterを使用します。

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.
    """
]

# クエリはシングルベクトルである必要があるため、チャンクを避けるためにchunk_sizeを-1に設定します。
# chunk_sizeが-1の場合、モデルはcls_vectorとすべてのトークン埋め込みの平均（mean_vector）からなる形状（2,2048）の配列を返します。
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]
# クエリベクトルはマルチベクトルを必要としないため、最終的なシングルベクトルとして平均のみを使用します
pred = [vecs[1:2, :] for vecs in query_vectors]

# spans_listは各チャンクのスパンを含み、スパンを使用してテキストを取得できます
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,  # fast_chunkがtrueの場合、入力ID上で直接チャンクし、それ以外の場合はRecursiveCharacterTextSplitterを使用します
)
# spans_listは各パッセージのスパンを保存し、passage_vectors_listは各パッセージのベクトルを保存するため、len(spans_list) == len(p_list)およびlen(spans_list) == len(passage_vectors_list)
# パッセージのスパンとベクトルについて、各スパンはベクトル（1*2048）に対応します。したがって、len(spans_list[idx]) ==  len(passage_vectors_list[idx])
print((query_vectors[0] @ passage_vectors_list[0].T).max())
# 出力 0.7331543
# 各チャンクの内容を取得
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()

このencodeの注釈を読んで、さらに情報を取得してください。

2.3.2 自分でテキストをチャンクする

自分でテキストをチャンクしたい場合は、encode関数のbatch_text_spansパラメータを設定するだけです。

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|>"

# モデルをロード
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

# テキストをチャンク
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はこのベクトルがclsベクトルであることを意味し、module_nameはcls_linearです。それ以外の場合、module_nameは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(
            # テキストの先頭に[CLS]トークンがあるため、1を追加します。
            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])
# 出力は (3, 2048) [0.01461297 0.02085092 0.0022509 ]

3 評価

3.1 MTEB(eng, v2)

URL: 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

3.2 LongEmbed

URL: http://mteb-leaderboard.hf.space/?benchmark_name=LongEmbed

再現スクリプト: https://huggingface.co/infgrad/dewey_en_beta/blob/main/scripts/evaluate/run_evaluate_long_embed.py

3.3 LoCoV1

URL: 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

結果: