Add support to run on Apple silicon using MPS

SeanLee97 / AnglE

Train and Infer Powerful Sentence Embeddings with AnglE | 🔥 SOTA on STS and MTEB Leaderboard

https://arxiv.org/abs/2309.12871

MIT License

398 stars 30 forks source link

Add support to run on Apple silicon using MPS #14

Closed quadrupole closed 7 months ago

quadrupole commented 7 months ago

When running AnglE on a Mac M1 computer, we are unable to move the model to the GPU because the repository assumes CUDA support. This PR makes some simple changes to allow for loading a model to create embeddings on a Mac M1.

Note I have not tested that this does not break when running on a CUDA enabled GPU.

quadrupole commented 7 months ago

@SeanLee97 Please let me know if this works and you are happy with the changes. Many thanks!

SeanLee97 commented 7 months ago

@quadrupole Looks good to me! Thanks for your contribution! I have bumped AnglE to v0.1.5: https://pypi.org/project/angle-emb/0.1.5/