jinaai

/

jina-colbert-v2

like 100

Follow

Jina AI 693

Hi @frederikschubert , yes, let me do it today or tomorrow

I'm also interested in using this for Vespa, let me see if I can export then quantize. Thanks for sharing this model btw -- arguably the top retriever model currently available :)

Hi @frederikschubert , the issue was due to the dynamic nature of the RoPE implementation. I've just uploaded the ONNX weights. Let me know how it works for you.

Great, thanks for your help! Would you mind sharing your conversion script? Either way, this is really helpful for us!

Sure, here it is:

import torch
from transformers import AutoModel, AutoTokenizer
import torch.onnx


model = AutoModel.from_pretrained('/home/admin/saba/jina-colbert-v2', trust_remote_code=True, use_flash_attn=False)
model.eval()

onnx_path =  '/home/admin/saba/jina-colbert-v2/onnx/model.onnx'

tokenizer = AutoTokenizer.from_pretrained('/home/admin/saba/jina-colbert-v2')
inputs = tokenizer(["jina", 'ai'], return_tensors="pt", padding='longest')
inps = inputs['input_ids']
mask = inputs['attention_mask']

torch.onnx.export(
    model,
    (inps, mask),
    onnx_path,
    export_params=True,
    do_constant_folding=True,
    input_names = ['input_ids', 'attention_mask'],
    output_names = ['text_embeds'],
    opset_version=16,
    dynamic_axes={
        'input_ids' : {0 : 'batch_size', 1: 'sequence_length'},
        'attention_mask' : {0 : 'batch_size', 1: 'sequence_length'},
        'text_embeds' : {0 : 'batch_size'}
    },
)

But the main challenge was to modify the model implementation to make it compatible with ONNX. That's why you were having the error.

Great, thanks again! :)

@jupyterjazz The exported onnx model in this repository works but the export script fails with the same error as above. I tried several ways of exporting the model but it seems that the RoPE embeddings still cause the export to fail.

I have the modified implementation that I used for conversion separately, it's not public atm. Can I ask why you want to convert the model yourself?

Alright! I am currently figuring out my general workflow of converting, optimizing and quantizing models and wanted to understand, why it fails for me. Would you mind sharing your diff? I can see why updating the public model for everyone might be a hurdle.

Okay, we can share it. Here's the modified implementation: https://huggingface.co/jinaai/xlm-roberta-flash-implementation-onnx

The README includes all the changes I made and explains the reasoning behind them, as well as the conversion script.

A few notes:

• The conversion script is written for the jina-embeddings-v3 model, which includes LoRA adapters. To make it work for ColBERT, you'll need to remove the task_id argument
• After conversion, it will probably output several files, each containing a part of the model. You can combine them using the convert_model_to_external_data function

Hope this helps!

Perfect! :) Thank you for the quick reply and for the code!