Improve documentation about what models can be used

[do-me]

Hi, thank you so much for open sourcing this repo. It's amazing and super useful to quickly get some insight into your data.
I created this Foursquare POI app for Italy with 3M points (GitHub) and it really helps me understand the data better! Operating at the limit of GitHub pages' free hosting with a 93Mb data file :D

I was just trying to figure out what kind of models could be used in embedding-atlas as I noticed you only mentioned the kind of dated all-MiniLM-L6-v2 in the API docs. I didn't quite dig in the code too much yet but at first glance it seems only the API version (and not the frontend) supports inferencing right?

So it would be great if you could:

• add some docs about what models work and what backend is used for inferencing in the API version (to understand whether it's already optimized for MPS for example)
• add support for frontend inferencing, e.g. with transformers.js. They already support WebGPU too. Just note that the batch size is absolutely crucial for speed, e.g. have a look at this demo I created to test batching. For reference: indexing the whole bible with 128 as batch size takes 35 seconds for 95605 chunks (M3 Max)

Batch size: 128
Chunks: 95605
Time passed: 35336.00 ms
Embeddings per second: 2705.60

Batch sizes of 64 or 256 instead roughly double the time needed on my device.

Apart from these standard models, I wanted to propose @MinishLab's static models (ideally with model2vec-rs) too as they are much faster on CPU and even rank higher on MTEB :)

[domoritz]

Super cool to hear that Embedding Atlas is useful for you and thanks for the link.

There are three ways to use Embedding atlas with increasingly more customization.

1. The web app which has no embedding model and umap in the browser for projection.
2. The cli which has a built in embedding model and map in Python for projection.
3. The jupyter/streamlit widget where you can use whatever code you want to do both the embedding and the projection.

I suppose this was unclear from the docs? Maybe this list needs to be at https://apple.github.io/embedding-atlas/overview.html. Would you be interested in sending a pull request with docs updates?

To unblock you, I recommend using the third option. If you have ideas for better embedding models or embedding models in the browser, we can consider those for enhancements.

[do-me]

I was already calculating the embeddings myself and passed them in the cli. However, there was a bit of a learning curve in the beginning to figure out what the app is capable of and/or doing automatically in which setting. Maybe something like a table could summarize it?

1. The web app which has no embedding model ...

Then I definitely see potential here for frontend-only inferencing, so that the user would only need to upload a parquet file with a text column! Thanks to WebGPU-support transformers.js is incredibly efficient. In August I'm on holiday but I might find some time in September to come up with a PR for this. I already coded a similar workflow (text -> chunking -> dim reduction -> map) in my own dim reduction component in SemanticFinder.

2. The cli which has a built in embedding model and map in Python for projection.

What backend are you using here? Plain pytorch, sentence-transformers, onnx...?

If you have ideas for better embedding models or embedding models in the browser, we can consider those for enhancements.

Absolutely, you can just have a look at https://huggingface.co/spaces/mteb/leaderboard. From my experience, bge-m3 works great for multilingual but is much biggert than allmini. Else, if you're into the smallest possible normal models, I'd consider gte-tiny or even gte-micro.
If I were to pick the fastest candidate, I'd probably go for Minish models via transformers.js (or via webassembly in the future MinishLab/model2vec-rs#25).
As I mentioned above, I could give it a shot in a PR, but probably in September :)

[domoritz]

What backend are you using here? Plain pytorch, sentence-transformers, onnx...?

embedding-atlas/packages/backend/embedding_atlas/cli.py

Line 122 in 0630e2f

help="Model name for generating embeddings (e.g., 'all-MiniLM-L6-v2').",

sets the model which then is used in

embedding-atlas/packages/backend/embedding_atlas/projection.py

Line 116 in 0630e2f

from sentence_transformers import SentenceTransformer

with SentenceTransformer. You can already pass in another model but it has to be supported by that code path.

[do-me]

Thanks for the quick feedback! Great, then static embedding models are already supported: https://huggingface.co/blog/static-embeddings

[domoritz]

I said something wrong. We already support using an embedding model if you upload your own data: https://apple.github.io/embedding-atlas/upload/.

[do-me]

Oh this is great so it's already supported!