search

lightonai / pylate

Late Interaction Models Training & Retrieval
https://lightonai.github.io/pylate/
MIT License
175 stars 7 forks source link

Add safetensor OR bin loading logic + add loading tests #73

Closed NohTow closed 3 days ago

NohTow commented 4 days ago

69 actually broke the loading of recent stanford-nlp models (small/jina), as they do not have bin.

That is on me, I checked the colbertv2 repository that had both safe tensor AND bin and rushed to the conclusion that all of them had both.

This PR merge old safe tensor loading and the new bin loading to be able to load both type of repository. We first try to load safe tensors and fall back to bin if not existing.

Also added some tests for loading different types of models (base, ST base, stanford-nlp bin/safetensor and PyLate), as well as a small check that should make sure that the model is not only loaded but correctly loaded.

raphaelsty commented 3 days ago

LGTM, in the next versions we might want to expand those tests