This is a Polish version of SPLADE++ (EnsembleDistil) model described in the paper From distillation to hard negative sampling: Making sparse neural ir models more effective. Sparse Lexical and Expansion (SPLADE) is a family of modern term-based retrieval methods employing Transformer language models. In this approach, the masked language modeling (MLM) head is optimized to generate a vocabulary-sized weight vector adapted for text retrieval. SPLADE is a highly effective sparse retrieval ranking algorithm, achieving results better than classic methods such as BM25 and comparable to high-quality dense encoders.
This model was fine-tuned from polish-distilroberta checkpoint on the Polish translation of the MS MARCO dataset. We used the default training hyperparameters from the official SPLADE repository.
Below is a example of using SPLADE without any additional dependencies other than Huggingface Transformers:
import torch, math
import numpy as np
from transformers import AutoTokenizer, AutoModelForMaskedLM
model_name = "sdadas/polish-splade"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)
vocab = {v: k for k, v in tokenizer.get_vocab().items()}
def encode_splade(text: str):
input = tokenizer([text], padding="longest", truncation=True, return_tensors="pt", max_length=512)
output = model(**input)
logits, attention_mask = output["logits"].detach(), input["attention_mask"].detach()
attention_mask = attention_mask.unsqueeze(-1)
vector = torch.max(torch.log(torch.add(torch.relu(logits), 1)) * attention_mask, dim=1)
vector = vector[0].detach().squeeze()
idx = np.nonzero(vector.cpu().numpy())[0]
vector = vector[idx]
return {vocab[k]: float(v) for k, v in zip(list(idx), list(vector))}
def cos_sim(vec1, vec2):
intersection = set(vec1.keys()) & set(vec2.keys())
numerator = sum([vec1[x] * vec2[x] for x in intersection])
sum1 = sum([vec1[x] ** 2 for x in list(vec1.keys())])
sum2 = sum([vec2[x] ** 2 for x in list(vec2.keys())])
denominator = math.sqrt(sum1) * math.sqrt(sum2)
return (numerator / denominator) if denominator else 0.0
question = encode_splade("Jak dożyć 100 lat?")
answer = encode_splade("Trzeba zdrowo się odżywiać i uprawiać sport.")
print(cos_sim(question, answer))
Example of use with the PIRB library:
from search import SpladeEncoder
from sentence_transformers.util import cos_sim
config = {"name": "sdadas/polish-splade", "fp16": True}
encoder = SpladeEncoder(config, True)
results = encoder.encode_batch(["Jak dożyć 100 lat?", "Trzeba zdrowo się odżywiać i uprawiać sport."])
print(cos_sim(results[0], results[1]))
Using SPLADE to index and search large datasets is a more complex task and requires integration with term-based index such as Lucene. For this purpose, you can use the official SPLADE implementation or reimplementation of this model in our PIRB library.
Citation
@article{dadas2024pirb,
title={{PIRB}: A Comprehensive Benchmark of Polish Dense and Hybrid Text Retrieval Methods},
author={Sławomir Dadas and Michał Perełkiewicz and Rafał Poświata},
year={2024},
eprint={2402.13350},
archivePrefix={arXiv},
primaryClass={cs.CL}
}