japanese-splade-base-v1

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

• 高性能な日本語SPLADE（スパース検索）モデルを公開しました
• SPLADE モデルの作り方・日本語SPLADEテクニカルレポート

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

利用方法

YASEM (Yet Another Splade|Sparse Embedder)

pip install yasem

from yasem import SpladeEmbedder

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

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

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

print(similarity)
# [[21.49299249 10.48868281  6.25582337]
#  [10.48868281 12.90587398  3.19429791]
#  [ 6.25582337  3.19429791 12.89678271]]

token_values = embedder.get_token_values(embeddings[0])

print(token_values)

#{
#    '車': 2.1796875,
#    '燃費': 2.146484375,
#    '向上': 1.7353515625,
#    '方法': 1.55859375,
#    '燃料': 1.3291015625,
#    '効果': 1.1376953125,
#    '良い': 0.873046875,
#    '改善': 0.8466796875,
#    'アップ': 0.833984375,
#    'いう': 0.70849609375,
#    '理由': 0.64453125,
# ...

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([[21.4943, 10.4816,  6.2540],
#         [10.4816, 12.9024,  3.1939],
#         [ 6.2540,  3.1939, 12.8919]])

ベンチマークスコア

retrieval (JMTEB)

JMTEB の評価結果です。japanese-splade-base-v1 は JMTEB をスパースベクトルで評価できるように変更したコードでの評価となっています。 なお、japanese-splade-base-v1 は jaqket, mrtydi のドメインを学習(testのデータ以外)しています。

reranking

JaCWIR

なお、japanese-splade-base-v1 は JaCWIR のドメインを学習していません。

JQaRA

なお、japanese-splade-base-v1 は JQaRA のドメイン(test以外)を学習したものとなっています。

学習元データセット

hpprc/emb から、auto-wiki-qa, mmarco, jsquad jaquad, auto-wiki-qa-nemotron, quiz-works quiz-no-mori, miracl, jqara mr-tydi, baobab-wiki-retrieval, mkqa データセットを利用しています。 また英語データセットとして、MS Marcoを利用しています。

注意事項

text-embeddings-inference で動かす場合、hotchpotch/japanese-splade-base-v1-dummy-fast-tokenizer-for-teiをご利用ください。