Rosetta Base Ja

🚀 RoSEtta

RoSEtta（RoFormerベースのSentence Encoder through Distillation）は、一般的な日本語テキスト埋め込みモデルで、検索タスクに特化しています。最大シーケンス長は1024で、長文の入力が可能です。CPU上でも実行でき、文章間の意味的な類似度を測定したり、クエリに基づいてパッセージを検索する検索システムとして機能するように設計されています。

🚀 クイックスタート

RoSEttaは、検索タスクに特化した一般的な日本語テキスト埋め込みモデルです。最大シーケンス長が1024で、長文の入力に対応しています。また、CPU上での実行も可能で、文章間の意味的な類似度を測定することができます。

✨ 主な機能

• RoPE（Rotary Position Embedding）を使用
• 最大シーケンス長が1024トークン
• 大規模な文章埋め込みモデルから知識蒸留
• 検索タスクに特化

推論時には、接頭辞 "query: " または "passage: " が必要です。詳細は「使用例」を参照してください。

📚 ドキュメント

モデルの説明

このモデルはRoFormerアーキテクチャに基づいています。MLM損失を使用して事前学習した後、弱教師付き学習を行いました。さらに、いくつかの大規模な埋め込みモデルを教師モデルとして知識蒸留と多段階の対照学習（GLuCoSE v2など）を行って学習を進めました。

使用例

基本的な使用法

SentenceTransformerを使用して推論を行うことができます。以下はサンプルコードです。

from sentence_transformers import SentenceTransformer
import torch.nn.functional as F

# Download from the 🤗 Hub
# The argument "trust_remote_code=True" is required to load the model
model = SentenceTransformer("pkshatech/RoSEtta-base-ja",trust_remote_code=True)

# Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
sentences = [
    'query: PKSHAはどんな会社ですか？',
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
    'query: 日本で一番高い山は？',
    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
]
embeddings = model.encode(sentences,convert_to_tensor=True)
print(embeddings.shape)
# [4, 768]

# Get the similarity scores for the embeddings
similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)
# [[1.0000, 0.5910, 0.4332, 0.5421],
# [0.5910, 1.0000, 0.4977, 0.6969],
# [0.4332, 0.4977, 1.0000, 0.7475],
# [0.5421, 0.6969, 0.7475, 1.0000]]

高度な使用法

Transformersを使用して推論を行うこともできます。以下はサンプルコードです。

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

def mean_pooling(last_hidden_states: Tensor,attention_mask: Tensor) -> Tensor:
    emb = last_hidden_states * attention_mask.unsqueeze(-1)
    emb = emb.sum(dim=1) / attention_mask.sum(dim=1).unsqueeze(-1)
    return emb

# Download from the 🤗 Hub
tokenizer = AutoTokenizer.from_pretrained("pkshatech/RoSEtta-base-ja")
# The argument "trust_remote_code=True" is required to load the model
model = AutoModel.from_pretrained("pkshatech/RoSEtta-base-ja",trust_remote_code=True)

# Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
sentences = [
    'query: PKSHAはどんな会社ですか？',
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
    'query: 日本で一番高い山は？',
    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
]

# Tokenize the input texts
batch_dict = tokenizer(sentences, max_length=1024, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = mean_pooling(outputs.last_hidden_state, batch_dict['attention_mask'])
print(embeddings.shape)
# [4, 768]

# Get the similarity scores for the embeddings
similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)
# [[1.0000, 0.5910, 0.4332, 0.5421],
# [0.5910, 1.0000, 0.4977, 0.6969],
# [0.4332, 0.4977, 1.0000, 0.7475],
# [0.5421, 0.6969, 0.7475, 1.0000]]

学習の詳細

RoSEttaの微調整は以下の手順で行われました。

ステップ1: 事前学習

• モデルはRoFormerアーキテクチャに基づいて事前学習されました。
• 学習データ: Japanese Wikipedia と cc100。

ステップ2: 弱教師付き学習

• 学習データ: MQA と mc4。

ステップ3: アンサンブル知識蒸留

• 埋め込み表現は、E5-mistral、gte-Qwen2、mE5-large を教師モデルとして知識蒸留されました。

ステップ4: 対照学習

• JSNLI、MNLI、PAWS-X、JSeM、Mr.TyDi からトリプレットを作成し、学習に使用しました。
• この学習は、文章埋め込みモデルとしての全体的な性能を向上させることを目的としています。

ステップ5: 検索特定の対照学習

• モデルを検索タスクに対してより頑健にするために、QAと検索タスクを用いた追加の2段階学習を行いました。
• 第1段階では、合成データセット auto-wiki-qa を使用して学習し、第2段階では JQaRA、MQA、Japanese Wikipedia Human Retrieval, Mr.TyDi,MIRACL, Quiz Works and Quiz No Mor を使用しました。

ベンチマーク

検索

MIRACL-ja、JQARA、JaCWIR、MLDR-ja を使用して評価しました。

注: JQARAとJaCWIRにおけるOpenAIの小規模埋め込みモデルの結果は、JQARA と JaCWIR から引用しています。

JMTEB

JMTEB を使用して評価しました。平均スコアはマクロ平均です。

著者

Chihiro Yano、Mocho Go、Hideyuki Tachibana、Hiroto Takegawa、Yotaro Watanabe

📄 ライセンス

このモデルは Apache License, Version 2.0 の下で公開されています。