from pprint import pprint

import torch
from sentence_transformers import SparseEncoder
from optimum.onnxruntime import ORTModelForMaskedLM

# 1. Load a pretrained SparseEncoder model
model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# Very hackishly override the model to use ORTModelForMaskedLM
class ORTModelForMaskedLMModule(ORTModelForMaskedLM, torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        torch.nn.Module.__init__(self)

model[0].auto_model = ORTModelForMaskedLMModule.from_pretrained("naver/splade-cocondenser-ensembledistil", library_name="transformers")

# Encode some texts into sparse embeddings
embeddings = model.encode(["I'm travelling to the grocery store to buy some milk."])

# Decode the embeddings again, the beauty of sparse embeddings
decoded = model.decode(embeddings, top_k=5)
pprint(decoded)
"""
[[('milk', 2.486112356185913),
  ('grocery', 1.7925951480865479),
  ('dairy', 1.6981983184814453),
  ('traveling', 1.5185797214508057),
  ('buy', 1.3063507080078125)]]
"""