Add utility for exporting a checkpoint to onnx format for serving in Vespa about ragatouille

[bclavie]

Hey Jo! Thank you for this, getting full Vespa compatibility is definitely something I want included very soon.

I don't think HF_ColBERT, in practice, does much more than VespaColBERT here in terms of model initialisation.

There's a potential clarity issue in inheriting BertPreTrainedModel: at the moment, all the open-source ColBERT models are BERT-based (and the attempts with RoBERta appear universally worse), but it might be an issue in the future if people initialise from better models/if a new, shinier encoder-only model appears and becomes the logical base rather than normal BERT. In practice, I don't think it'll actually hinder anything, since you're only using components that you're manually loading/overwriting.

The bert and linear property should both be easy to export, since RAGatouille uses the normal ColBERT training code, the model that gets trained will have them as they're defined features of BaseColBERT. The bert property here is a bit of misnomer, it actually stores any underlying model, even if not BERT 🤔

JaColBERT was trained with RAGatouille-like utils (although quite different back then!), and this works as expected (all models have similar, properly-initiated weights):

VespaJaColBERT = VespaColBERT.from_pretrained("bclavie/JaColBERT", dim=128)
HFColBERT = class_factory("bclavie/JaColBERT")
HFJaColBERT = HFColBERT.from_pretrained("bclavie/JaColBERT", colbert_config=ColBERTConfig(dim=128))
BaseColBERT = BaseColBERT("bclavie/JaColBERT")

assert torch.equal(HFJaColBERT.linear.weight, VespaJaColBERT.linear.weight) and \
       torch.equal(HFJaColBERT.linear.weight, BaseColBERT.linear.weight)

assert torch.equal(HFJaColBERT.bert.embeddings.word_embeddings.weight, VespaJaColBERT.bert.embeddings.word_embeddings.weight) and \
         torch.equal(HFJaColBERT.bert.embeddings.word_embeddings.weight, BaseColBERT.bert.embeddings.word_embeddings.weight)

# assert torch.equal(VespaJaColBERT.linear.weight, VespaColBERTV2.linear.weight) crashes as expected, too!

I think that adding your current exporter, mostly as-is (if you're fine with me re-using the code!) would be fine, and have a ragatouille.models.export_to_vespa function which could be called via RAGTrainer and RAGPreTrainedModel as well would make sense.

I'm also planning on adding an export_to_huggingface_hub() soon (requiring a credential as a env variable), and will make sure it has an export_vespa_onnx: bool parameter to it 😄

[jobergum]

Thanks for that! having this type of utility would make it a lot easier to promote RAGatouille as a way to train ColBERT models!

I think that adding your current exporter, mostly as-is (if you're fine with me re-using the code!) would be fine, and have a ragatouille.models.export_to_vespa function which could be called via RAGTrainer and RAGPreTrainedModel as well would make sense.

Fantastic, would greatly appreciate it! I'm going to try to open a PR to upload onnx artifacts to colbert-ir/colbertv2.0, cc @okhat

[bclavie]

(tentative implementation in #19)

[bclavie]

Implemented in #19