High-performance embeddings with Apple MLX π
Version 0.2.0
llamamlx-embeddings is a Python library that provides high-performance text embeddings using Apple's MLX framework, optimized for Apple Silicon. It offers a unified interface for generating embeddings with various models, efficient batch processing, quantization options, seamless integration with vector databases, and easy deployment as a FastAPI service.
This package follows a standardized structure for ease of use and maintainability:
llamamlx-embeddings/
βββ src/ # Source code directory
β βββ llamamlx_embeddings/ # Main package
β βββ api/ # API interfaces and handlers
β βββ benchmarks/ # Benchmarking tools
β βββ core/ # Core functionality
β βββ conversion/ # Model conversion utilities
β βββ integrations/ # Vector DB integrations
β βββ processing/ # Text processing utilities
β βββ quantization/ # Model quantization tools
β βββ utils/ # Common utility functions
β βββ visualization/ # Visualization utilities
β βββ __init__.py # Package initialization
β βββ cli.py # Command-line interface
β βββ client.py # API client
β βββ logging.py # Logging configuration
β βββ version.py # Version information
βββ tests/ # Test directory
βββ docs/ # Documentation
βββ examples/ # Example scripts
βββ benchmarks/ # Benchmark results
βββ setup.py # Package setup script
βββ pyproject.toml # Project configuration
βββ MANIFEST.in # Package manifest
βββ README.md # Project README
βββ LICENSE # License information
On Apple M2 Pro, using batch size 32:
| Model | Texts/sec | Dim | Type |
|---|---|---|---|
| BAAI/bge-small-en-v1.5 | ~245 | 384 | Dense |
| sentence-transformers/all-MiniLM-L6-v2 | ~285 | 384 | Dense |
| intfloat/e5-small-v2 | ~230 | 384 | Dense |
| prithivida/Splade_PP_en_v1 | ~80 | var | Sparse |
With INT8 quantization, throughput improves by ~30% and model size reduces by ~69%
# Basic installation
pip install llamamlx-embeddings
# With vector database integrations
pip install llamamlx-embeddings[qdrant,pinecone]
# Full installation with all features
pip install llamamlx-embeddings[all]
git clone https://github.com/yourusername/llamamlx-embeddings.git
cd llamamlx-embeddings
pip install -e .
from llamamlx_embeddings import TextEmbedding
import numpy as np
# Create an embedding model (will download if needed)
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Generate embeddings
query = "How to make a delicious pizza?"
query_embedding = model.embed_query(query)
documents = [
"Pizza is a dish of Italian origin consisting of a usually round, flat base of leavened wheat-based dough.",
"To make pizza, you need flour, water, yeast, salt, olive oil, tomato sauce, and cheese."
]
doc_embeddings = model.embed_documents(documents)
# Calculate similarities
for i, doc_emb in enumerate(doc_embeddings):
similarity = np.dot(query_embedding, doc_emb) / (np.linalg.norm(query_embedding) * np.linalg.norm(doc_emb))
print(f"Document {i+1} similarity: {similarity:.4f}")
from llamamlx_embeddings import MockEmbedding
# Create a mock embedding model
model = MockEmbedding(dimensions=384)
# Use it like a regular embedding model
query_embedding = model.embed_query("How to make pizza?")
document_embeddings = model.embed_documents(["Document 1", "Document 2"])
# Perfect for testing applications without downloading large models
Start the server:
llamamlx-embeddings serve --host 0.0.0.0 --port 8000
Use the client:
from llamamlx_embeddings import LlamamlxEmbeddingsClient
# Create a client
client = LlamamlxEmbeddingsClient(base_url="http://localhost:8000")
# Generate embeddings
query = "How to make a delicious pizza?"
query_embedding = client.get_embeddings(query, is_query=True)[0]
documents = [
"Pizza is a dish of Italian origin consisting of a usually round, flat base of leavened wheat-based dough.",
"To make pizza, you need flour, water, yeast, salt, olive oil, tomato sauce, and cheese."
]
doc_embeddings = client.get_embeddings(documents)
For comprehensive documentation, visit our documentation site.
Dense models:
Sparse models:
Late interaction models:
Cross-encoder models:
from llamamlx_embeddings import TextEmbedding, QdrantClient
# Create embedding model
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Initialize Qdrant client
vector_db = QdrantClient(
url="https://your-qdrant-instance.com",
collection_name="my_collection",
embedding_model=model
)
# Add documents
vector_db.add(
documents=["Document 1 text", "Document 2 text"],
metadata=[{"source": "file1.txt"}, {"source": "file2.txt"}]
)
# Search with query
results = vector_db.query("My search query", limit=5)
from llamamlx_embeddings import TextEmbedding
# Load a quantized model
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", quantize=True)
# Query and document embeddings work the same way
query_embedding = model.embed_query("How to make pizza?")
from llamamlx_embeddings import add_custom_model, TextEmbedding
# Add a custom model
add_custom_model(
model_name="my-custom-model",
model_path="/path/to/model/files",
model_type="dense",
dimensions=768,
description="My custom embedding model"
)
# Use the custom model
model = TextEmbedding(model_name="my-custom-model")
Contributions are welcome! Please check out our contributing guide to get started.
This project is licensed under the MIT License - see the LICENSE file for details.
Python
99.2%
High-performance embeddings with Apple MLX π
Version 0.2.0
llamamlx-embeddings is a Python library that provides high-performance text embeddings using Apple's MLX framework, optimized for Apple Silicon. It offers a unified interface for generating embeddings with various models, efficient batch processing, quantization options, seamless integration with vector databases, and easy deployment as a FastAPI service.
This package follows a standardized structure for ease of use and maintainability:
llamamlx-embeddings/
βββ src/ # Source code directory
β βββ llamamlx_embeddings/ # Main package
β βββ api/ # API interfaces and handlers
β βββ benchmarks/ # Benchmarking tools
β βββ core/ # Core functionality
β βββ conversion/ # Model conversion utilities
β βββ integrations/ # Vector DB integrations
β βββ processing/ # Text processing utilities
β βββ quantization/ # Model quantization tools
β βββ utils/ # Common utility functions
β βββ visualization/ # Visualization utilities
β βββ __init__.py # Package initialization
β βββ cli.py # Command-line interface
β βββ client.py # API client
β βββ logging.py # Logging configuration
β βββ version.py # Version information
βββ tests/ # Test directory
βββ docs/ # Documentation
βββ examples/ # Example scripts
βββ benchmarks/ # Benchmark results
βββ setup.py # Package setup script
βββ pyproject.toml # Project configuration
βββ MANIFEST.in # Package manifest
βββ README.md # Project README
βββ LICENSE # License information
On Apple M2 Pro, using batch size 32:
| Model | Texts/sec | Dim | Type |
|---|---|---|---|
| BAAI/bge-small-en-v1.5 | ~245 | 384 | Dense |
| sentence-transformers/all-MiniLM-L6-v2 | ~285 | 384 | Dense |
| intfloat/e5-small-v2 | ~230 | 384 | Dense |
| prithivida/Splade_PP_en_v1 | ~80 | var | Sparse |
With INT8 quantization, throughput improves by ~30% and model size reduces by ~69%
# Basic installation
pip install llamamlx-embeddings
# With vector database integrations
pip install llamamlx-embeddings[qdrant,pinecone]
# Full installation with all features
pip install llamamlx-embeddings[all]
git clone https://github.com/yourusername/llamamlx-embeddings.git
cd llamamlx-embeddings
pip install -e .
from llamamlx_embeddings import TextEmbedding
import numpy as np
# Create an embedding model (will download if needed)
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Generate embeddings
query = "How to make a delicious pizza?"
query_embedding = model.embed_query(query)
documents = [
"Pizza is a dish of Italian origin consisting of a usually round, flat base of leavened wheat-based dough.",
"To make pizza, you need flour, water, yeast, salt, olive oil, tomato sauce, and cheese."
]
doc_embeddings = model.embed_documents(documents)
# Calculate similarities
for i, doc_emb in enumerate(doc_embeddings):
similarity = np.dot(query_embedding, doc_emb) / (np.linalg.norm(query_embedding) * np.linalg.norm(doc_emb))
print(f"Document {i+1} similarity: {similarity:.4f}")
from llamamlx_embeddings import MockEmbedding
# Create a mock embedding model
model = MockEmbedding(dimensions=384)
# Use it like a regular embedding model
query_embedding = model.embed_query("How to make pizza?")
document_embeddings = model.embed_documents(["Document 1", "Document 2"])
# Perfect for testing applications without downloading large models
Start the server:
llamamlx-embeddings serve --host 0.0.0.0 --port 8000
Use the client:
from llamamlx_embeddings import LlamamlxEmbeddingsClient
# Create a client
client = LlamamlxEmbeddingsClient(base_url="http://localhost:8000")
# Generate embeddings
query = "How to make a delicious pizza?"
query_embedding = client.get_embeddings(query, is_query=True)[0]
documents = [
"Pizza is a dish of Italian origin consisting of a usually round, flat base of leavened wheat-based dough.",
"To make pizza, you need flour, water, yeast, salt, olive oil, tomato sauce, and cheese."
]
doc_embeddings = client.get_embeddings(documents)
For comprehensive documentation, visit our documentation site.
Dense models:
Sparse models:
Late interaction models:
Cross-encoder models:
from llamamlx_embeddings import TextEmbedding, QdrantClient
# Create embedding model
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Initialize Qdrant client
vector_db = QdrantClient(
url="https://your-qdrant-instance.com",
collection_name="my_collection",
embedding_model=model
)
# Add documents
vector_db.add(
documents=["Document 1 text", "Document 2 text"],
metadata=[{"source": "file1.txt"}, {"source": "file2.txt"}]
)
# Search with query
results = vector_db.query("My search query", limit=5)
from llamamlx_embeddings import TextEmbedding
# Load a quantized model
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", quantize=True)
# Query and document embeddings work the same way
query_embedding = model.embed_query("How to make pizza?")
from llamamlx_embeddings import add_custom_model, TextEmbedding
# Add a custom model
add_custom_model(
model_name="my-custom-model",
model_path="/path/to/model/files",
model_type="dense",
dimensions=768,
description="My custom embedding model"
)
# Use the custom model
model = TextEmbedding(model_name="my-custom-model")
Contributions are welcome! Please check out our contributing guide to get started.
This project is licensed under the MIT License - see the LICENSE file for details.
Python
99.2%