Tutorial to export a SPLADE model to ONNX about splade HOT 6 OPEN

import torch
from transformers import AutoModelForMaskedLM,AutoTokenizer # type: ignore

class TransformerRep(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = AutoModelForMaskedLM.from_pretrained("naver/splade-cocondenser-ensembledistil", torchscript=True) # type: ignore
        self.model.eval() # type: ignore
        self.fp16 = True

    def encode(self, input_ids, token_type_ids, attention_mask):
        # Tokens is a dict with keys input_ids and attention_mask
        return self.model(input_ids, token_type_ids, attention_mask)[0]



class SpladeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = TransformerRep()
        self.agg = "max"
        self.model.eval()
    
    def forward(self, input_ids,token_type_ids, attention_mask):
        with torch.cuda.amp.autocast(): # type: ignore
            with torch.no_grad():
                lm_logits = self.model.encode(input_ids,token_type_ids, attention_mask)[0]
                vec, _ = torch.max(torch.log(1 + torch.relu(lm_logits)) * attention_mask.unsqueeze(-1), dim=1)
                indices = vec.nonzero().squeeze()
                weights = vec.squeeze()[indices]
        return indices[:,1], weights[:,1]

# Convert the model to TorchScript
model = SpladeModel()
tokenizer = AutoTokenizer.from_pretrained("naver/splade-cocondenser-ensembledistil")
sample = "the capital of france is paris"
inputs = tokenizer(sample, return_tensors="pt")
traced_model = torch.jit.trace(model, (inputs["input_ids"], inputs["token_type_ids"], inputs["attention_mask"]))

import torch
dyn_axis = {
    'input_ids': {0: 'batch_size', 1: 'sequence'},
    'attention_mask': {0: 'batch_size', 1: 'sequence'},
    'token_type_ids': {0: 'batch_size', 1: 'sequence'},
    'indices': {0: 'batch_size', 1: 'sequence'},
    'weights': {0: 'batch_size', 1: 'sequence'}
    }
model = torch.jit.load(model_file)
onnx_model = torch.onnx.export(
    model,
    dummy_input, # type: ignore
    f=model_onnx_file,
    input_names=['input_ids','token_type_ids', 'attention_mask'],
    output_names=['indices', 'weights'],
    dynamic_axes=dyn_axis,
    do_constant_folding=True,
    opset_version=15,
    verbose=False,
)

model_names= [
   "naver/splade_v2_max",
   "naver/splade_v2_distil",
   "naver/splade-cocondenser-ensembledistil",
   "naver/efficient-splade-VI-BT-large-query",
   "naver/efficient-splade-VI-BT-large-doc",
]