polish-splade

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))

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]))