castorini / ura-projects

0 stars 1 forks source link

Pyserini: Port Torch Models to MLX and enable Indexing and Searching on Mac M series machines #39

Open ToluClassics opened 2 months ago

ToluClassics commented 2 months ago

Summary: The goal is to extend the encode module (https://github.com/castorini/pyserini/tree/master/pyserini/encode ) of Pyserini to allow loading and inferencing retriever models in MLX similar to the way it is currently done in PyTorch. This enables us to index and search effectively on Mac M series laptops with near Cuda speed. So similar to DprDocumentEncoder we will have something like MlxDprDocumentEncoder.

What this entails is translating model architecture code from PyTorch to MLX. The interface of both frameworks are very similar except that MLX is more numpy like