USER Base

🚀 USER-base

USER-base（ロシア語向けの汎用文書エンコーダー、Universal Sentence Encoder for Russian）は、sentence-transformer をベースにしたモデルで、ロシア語の文章から埋め込みベクトルを抽出することを専門としています。このモデルは、文章や段落を 768 次元の密なベクトル空間にマッピングし、クラスタリングや意味検索などのタスクに利用できます。このモデルは、deepvk/deberta-v1-base を初期化の基盤とし、ロシア語に特化して訓練されており、他の言語での性能は評価されていません。

🚀 クイックスタート

このモデルを使用するには、sentence-transformers をインストールする必要があります。

pip install -U sentence-transformers

次に、以下のようにモデルを使用できます。

from sentence_transformers import SentenceTransformer

queries = [
  "Когда был спущен на воду первый миноносец «Спокойный»?",
  "Есть ли нефть в Удмуртии?"
]
passages = [
  "Спокойный (эсминец)\nЗачислен в списки ВМФ СССР 19 августа 1952 года.",
  "Нефтепоисковые работы в Удмуртии были начаты сразу после Второй мировой войны в 1945 году и продолжаются по сей день. Добыча нефти началась в 1967 году."
]

model = SentenceTransformer("deepvk/USER-base")
# Prompt should be specified according to the task (either 'query' or 'passage').
passage_embeddings = model.encode(passages, normalize_embeddings=True, prompt_name='passage')
# For tasks other than retrieval, you can simply use the `query` prompt, which is set by default.
query_embeddings = model.encode(queries, normalize_embeddings=True)

ただし、transformers を直接使用してモデルを呼び出すこともできます。

import torch.nn.functional as F
from torch import Tensor, inference_mode
from transformers import AutoTokenizer, AutoModel

def average_pool(
  last_hidden_states: Tensor,
  attention_mask: Tensor
) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(
      ~attention_mask[..., None].bool(), 0.0
    )
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

# You should manually add prompts when using the model directly. Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
input_texts = [
  "query: Когда был спущен на воду первый миноносец «Спокойный»?",
  "query: Есть ли нефть в Удмуртии?",
  "passage: Спокойный (эсминец)\nЗачислен в списки ВМФ СССР 19 августа 1952 года.",
  "passage: Нефтепоисковые работы в Удмуртии были начаты сразу после Второй мировой войны в 1945 году и продолжаются по сей день. Добыча нефти началась в 1967 году."
]

tokenizer = AutoTokenizer.from_pretrained("deepvk/USER-base")
model = AutoModel.from_pretrained("deepvk/USER-base")

batch_dict = tokenizer(
  input_texts, padding=True, truncation=True, return_tensors="pt"
)
with inference_mode():
  outputs = model(**batch_dict)
  embeddings = average_pool(
    outputs.last_hidden_state, batch_dict["attention_mask"]
  )
  embeddings = F.normalize(embeddings, p=2, dim=1)

# Scores for query-passage
scores = (embeddings[:2] @ embeddings[2:].T) * 100
# [[55.86, 30.95],
#  [22.82, 59.46]]
print(scores.round(decimals=2))

⚠️ 重要な注意事項

各入力テキストは "query: " または "passage: " で始める必要があります。検索以外のタスクでは、"query: " プレフィックスを直接使用できます。

✨ 主な機能

• ロシア語に特化して設計されており、ロシア語の文章から効果的に埋め込みベクトルを抽出できます。
• 文章や段落を 768 次元の密なベクトル空間にマッピングし、クラスタリングや意味検索などのタスクに適しています。

📦 インストール

このモデルを使用するには、sentence-transformers をインストールする必要があります。

pip install -U sentence-transformers

💻 使用例

基本的な使用法

from sentence_transformers import SentenceTransformer

queries = [
  "Когда был спущен на воду первый миноносец «Спокойный»?",
  "Есть ли нефть в Удмуртии?"
]
passages = [
  "Спокойный (эсминец)\nЗачислен в списки ВМФ СССР 19 августа 1952 года.",
  "Нефтепоисковые работы в Удмуртии были начаты сразу после Второй мировой войны в 1945 году и продолжаются по сей день. Добыча нефти началась в 1967 году."
]

model = SentenceTransformer("deepvk/USER-base")
# Prompt should be specified according to the task (either 'query' or 'passage').
passage_embeddings = model.encode(passages, normalize_embeddings=True, prompt_name='passage')
# For tasks other than retrieval, you can simply use the `query` prompt, which is set by default.
query_embeddings = model.encode(queries, normalize_embeddings=True)

高度な使用法

import torch.nn.functional as F
from torch import Tensor, inference_mode
from transformers import AutoTokenizer, AutoModel

def average_pool(
  last_hidden_states: Tensor,
  attention_mask: Tensor
) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(
      ~attention_mask[..., None].bool(), 0.0
    )
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

# You should manually add prompts when using the model directly. Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
input_texts = [
  "query: Когда был спущен на воду первый миноносец «Спокойный»?",
  "query: Есть ли нефть в Удмуртии?",
  "passage: Спокойный (эсминец)\nЗачислен в списки ВМФ СССР 19 августа 1952 года.",
  "passage: Нефтепоисковые работы в Удмуртии были начаты сразу после Второй мировой войны в 1945 году и продолжаются по сей день. Добыча нефти началась в 1967 году."
]

tokenizer = AutoTokenizer.from_pretrained("deepvk/USER-base")
model = AutoModel.from_pretrained("deepvk/USER-base")

batch_dict = tokenizer(
  input_texts, padding=True, truncation=True, return_tensors="pt"
)
with inference_mode():
  outputs = model(**batch_dict)
  embeddings = average_pool(
    outputs.last_hidden_state, batch_dict["attention_mask"]
  )
  embeddings = F.normalize(embeddings, p=2, dim=1)

# Scores for query-passage
scores = (embeddings[:2] @ embeddings[2:].T) * 100
# [[55.86, 30.95],
#  [22.82, 59.46]]
print(scores.round(decimals=2))

🔧 技術詳細

訓練の詳細

このモデルの訓練は、bge-base-en の訓練アルゴリズムを基に、さらに改良されています。

• 初期化：deepvk/deberta-v1-base を基に初期化されます。
• 第 1 段階：mMarco コーパス のロシア語部分で弱教師付きの対比事前訓練を行います。
• 第 2 段階：データの対称性に基づいて 2 つの異なるモデルを教師付きで微調整し、LM-Cocktail を用いて統合します。
  1. 命令設計を変更し、多言語の方法を簡素化して推論を容易にします。対称データ (S1, S2) には、命令 "query: S1" と "query: S2" を使用します。非対称データには、"query: S1" と "passage: S2" を使用します。
  2. データを分割したため、対称モデルに AnglE 損失 を追加で適用し、対称タスクの性能を向上させます。
  3. 最後に、2 つのモデルを統合し、LM-Cocktail で統合重みを調整して、最終的なモデル USER を生成します。

データセット

モデル開発の過程で、2 つのデータセットを追加で収集しました：deepvk/ru-HNP と deepvk/ru-WANLI。

総正例ペア：3,352,653
総負例ペア：792,644（AIINLI、MIRACL、deepvk/ru-WANLI、deepvk/ru-HNP の負例ペア）

すべてのラベル付きデータセットでは、微調整に訓練セットのみを使用します。Gazeta、Mlsum、Xlsum データセットでは、(タイトル/テキスト) および (タイトル/要約) のペアを組み合わせて、非対称データとして使用します。

AllNLI は、SNLI、MNLI、ANLI のロシア語翻訳の組み合わせです。

実験

• ベンチマークの選択：encodechka ランキングの現在のトップモデルをベースラインとして選択します。また、MTEB のロシア語サブセットでモデルを評価します。このサブセットには 10 のタスクが含まれています。計算リソースの制限により、一部の MTEB タスク（特にクラスタリングタスク）では bge-m3 を検証できませんでした。これら 2 つのベンチマークに加えて、MIRACL でもモデルを評価しました。
• 実験環境：すべての実験は、NVIDIA TESLA A100 40 GB GPU を使用して行いました。各タスクの公式リポジトリの検証スクリプトを使用しました。

モデルのサイズが表示されており、大きなモデルは他のモデルと視覚的に区別されます。指標の絶対的なリーダーは太字で強調され、同じ規模のモデルの中でのリーダーは下線が引かれています。

このモデルは、Encodechka と MTEB の両方で、同じ規模の他のすべてのモデルを上回っています。ただし、検索タスクでは既存の解決策にやや劣っているため、将来的な研究でこの問題を解決する予定です。

📚 ドキュメント

よくある質問

入力テキストに "query: " と "passage: " のプレフィックスを追加する必要がありますか？ はい、モデルはこのように訓練されているため、追加しないと性能が低下します。以下はいくつかの経験則です。

• 非対称タスク（オープン質問応答の段落検索、一時的な情報検索など）では、それぞれ "query: " と "passage: " を使用します。
• 対称タスク（意味的な類似性、バイリンガルマイニング、言い換え検索など）では、"query: " プレフィックスを使用します。
• 埋め込みベクトルを特徴として使用する場合（線形検出分類、クラスタリングなど）、"query: " プレフィックスを使用します。

引用

@misc{deepvk2024user,
    title={USER: Universal Sentence Encoder for Russian},
    author={Malashenko, Boris and  Zemerov, Anton and Spirin, Egor},
    url={https://huggingface.co/datasets/deepvk/USER-base},
    publisher={Hugging Face}
    year={2024},
}

📄 ライセンス

このプロジェクトは、Apache-2.0 ライセンスの下で提供されています。