高性能な日本語 SPLADE (Sparse Lexical and Expansion Model) モデルです。テキストからスパースベクトルへの変換デモで、どのようなスパースベクトルに変換できるか、WebUI から気軽にお試しいただけます。

• 情報検索モデルで最高性能(512トークン以下)・日本語版SPLADE v2をリリース

また、モデルの学習にはYAST - Yet Another SPLADE or Sparse Trainerを使っています。

利用方法

YASEM (Yet Another Splade|Sparse Embedder)

YASEM を利用することで、SPLADEの推論・単語トークンの確認を簡単に行えます。

pip install yasem

from yasem import SpladeEmbedder

model_name = "hotchpotch/japanese-splade-v2"
embedder = SpladeEmbedder(model_name)

query = "車の燃費を向上させる方法は？"
docs = [
    "急発進や急ブレーキを避け、一定速度で走行することで燃費が良くなります。",
    "車の運転時、急発進や急ブレーキをすると、燃費が悪くなります。",
    "車を長持ちさせるには、消耗品を適切なタイミングで交換することが重要です。",
]

print(embedder.rank(query, docs, return_documents=True))

[
 { 'corpus_id': 0
 , 'score': 4.28
 , 'text': '急発進や急ブレーキを避け、一定速度で走行することで燃費が良くなります。' }
 ,
 { 'corpus_id': 2
 , 'score': 2.47
 , 'text': '車を長持ちさせるには、消耗品を適切なタイミングで交換することが重要です。' }
 ,
 { 'corpus_id': 1
 , 'score': 2.34
 , 'text': '車の運転時、急発進や急ブレーキをすると、燃費が悪くなります。' }
]

sentences = [query] + docs

embeddings = embedder.encode(sentences)
similarity = embedder.similarity(embeddings, embeddings)

print(similarity)

[[5.19151189, 4.28027662, 2.34164901, 2.47221905],
[4.28027662, 11.64426784, 5.00328318, 2.15031016],
[2.34164901, 5.00328318, 6.05594296, 1.33752085],
[2.47221905, 2.15031016, 1.33752085, 9.39414744]]

token_values = embedder.get_token_values(embeddings[0])
print(token_values)

{
 '燃費': 1.13,
 '方法': 1.07,
 '車': 1.05,
 '高める': 0.67,
 '向上': 0.56,
 '増加': 0.52,
 '都市': 0.44,
 'ガソリン': 0.32,
 '改善': 0.30,
 ...

transformers からの利用

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch

model = AutoModelForMaskedLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def splade_max_pooling(logits, attention_mask):
    relu_log = torch.log(1 + torch.relu(logits))
    weighted_log = relu_log * attention_mask.unsqueeze(-1)
    max_val, _ = torch.max(weighted_log, dim=1)
    return max_val

tokens = tokenizer(
    sentences, return_tensors="pt", padding=True, truncation=True, max_length=512
)
tokens = {k: v.to(model.device) for k, v in tokens.items()}

with torch.no_grad():
    outputs = model(**tokens)
embeddings = splade_max_pooling(outputs.logits, tokens["attention_mask"])

similarity = torch.matmul(embeddings.unsqueeze(0), embeddings.T).squeeze(0)
print(similarity)

tensor([
   [5.1872, 4.2792, 2.3440, 2.4680],
   [4.2792, 11.6327, 4.9983, 2.1470],
   [2.3440, 4.9983, 6.0517, 1.3377],
   [2.4680, 2.1470, 1.3377, 9.3801]
])

ベンチマークスコア

retrieval (JMTEB)

JMTEB の評価結果です。

japanese-splade-v2 は JMTEB をスパースベクトルで評価できるように変更したコードでの評価となっています。 なお、japanese-splade-v2 は JMTEB タスクである jaqket(や派生のjaqra), mrtydi(と派生のmiracl), jagovfaqs, nlp_jornal のデータセットのtrain,dev, testなどのデータは 学習に利用していません。

スパース性

v1 ではスパース性が強すぎたので、v2 ではバランスをとったスパース性を持たせています。

• https://github.com/hotchpotch/yast/blob/main/utils/JMTEB_L0.py

で計測しています。

学習元データセット

• hpprc/emb
  • auto-wiki-qa
  • jsquad
  • jaquad
  • auto-wiki-qa-nemotron
  • quiz-works
  • quiz-no-mori
  • baobab-wiki-retrieval
  • mkqa
• hotchpotch/japanese-splade-v1-hard-negatives
  • mmarco
  • mqa
  • msmarco-ja
• hotchpotch/mmarco-hard-negatives-reranker-score
  • english

ライセンス

MIT