Add embedder-tool CLI for document indexing and semantic search

[rudrankriyam]

This PR introduces a new command-line tool, embedder-tool, that helps to work with embedding models from the MLXEmbedders library. The tool can index documents, semantic search, and for model management.

Features

index

• Create embedding indexes from document corpora
• Supports multiple file formats (txt, md, etc.)
• Recursive directory scanning
• Batch processing with configurable batch sizes
• Progress reporting during indexing
• Handles binary detection and encoding errors
• Outputs JSON index files with document paths and embeddings

search

• Semantic search over indexed documents
• Cosine similarity-based ranking
• Configurable result count
• Dimension mismatch detection and reporting
• Returns ranked results with similarity scores

list

• List registered embedding model configurations
• Shows all available pre-configured models
• Optional display of local directory models
• Sorted output for easy browsing

demo

• Indexes sample documentation from the repository
• Runs example queries showing semantic search

Bug Fixes

For NomicBert Model Architecture. Let me know if this should be a separate PR.

• Fixed pooler Linear layer to use bias: false matching model weights
• Updated model creation to use pooler: false for nomic-embed-text-v1.5

Usage Examples

Index documents

embedder-tool index --directory docs/ --output index.json --extensions md txt --recursive

Search the index

embedder-tool search --index index.json --query "How do embeddings work?" --top 5

List available models

embedder-tool list

Run demo

embedder-tool demo

Closes #407

[davidkoski]

This is really cool! The demo command worked well. When I tried to build an index on my own:

./mlx-run --debug embedder-tool index --directory /Users/dkoski/Developer/mlx-swift-examples --recursive --output /tmp/index --normalize

I hit this -- it was doing an eval of the graph:

MLX/ErrorHandler.swift:332: Fatal error: [metal::malloc] Attempting to allocate 76956753216 bytes which is greater than the maximum allowed buffer size of 41747087360 bytes.

I found that if I set the batch-size to 4 it ran OK, which is what it looks like the demo command does internally. Should that be the default? Are there situations where 32 will work?

It also seems to use a lot of memory -- I wonder if doing something like GPU.set(cacheLimit: 1024 * 1024 * 25) or using MemoryArguments from llm-tool would be helpful? I don't think this is harmful as it uses memory the you have, but it may be misleading if somebody looks at the memory use.

[rudrankriyam]

Yep, exactly. I also added a repl command which is more fun to play haha. Looking into the memory right now that's why converted to draft after I hit it

[rudrankriyam]

Back to this PR; let's get it done!

[davidkoski]

It looks like this just needs swift-format and then we can gt it merged. Thanks!

[rudrankriyam]

Done!

[davidkoski]

Thank you, this is a big help!