333 lines
7.8 KiB
Markdown
333 lines
7.8 KiB
Markdown
# Configurable Embeddings Models
|
|||
|
|
|
||
|
|
**Status**: Implemented
|
||
|
|
**Feature**: Customer-selectable embedding models via `EMBEDDINGS_MODEL` environment variable
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Overview
|
||
|
|
|
||
|
|
The memory system supports multiple embedding models, all configured to return **exactly 768 dimensions** to match the pgvector schema and HNSW indexes.
|
||
|
|
|
||
|
|
Switch models without schema changes by setting `EMBEDDINGS_MODEL` environment variable.
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Supported Models
|
||
|
|
|
||
|
|
### 1. nomic-ai/nomic-embed-text-v2-moe (Default)
|
||
|
|
|
||
|
|
**Characteristics**:
|
||
|
|
- **Dimensions**: 768
|
||
|
|
- **Speed**: Fast (MoE optimization)
|
||
|
|
- **Quality**: Good
|
||
|
|
- **Languages**: 30+ (multilingual)
|
||
|
|
- **Use case**: Default, production recommended
|
||
|
|
|
||
|
|
```bash
|
||
|
|
export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v2-moe
|
||
|
|
```
|
||
|
|
|
||
|
|
### 2. nomic-ai/nomic-embed-text-v1.5
|
||
|
|
|
||
|
|
**Characteristics**:
|
||
|
|
- **Dimensions**: 768
|
||
|
|
- **Speed**: Slower than v2-moe
|
||
|
|
- **Quality**: Slightly better than v2-moe
|
||
|
|
- **Languages**: 30+ (multilingual)
|
||
|
|
- **Use case**: When quality matters more than speed
|
||
|
|
|
||
|
|
```bash
|
||
|
|
export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v1.5
|
||
|
|
```
|
||
|
|
|
||
|
|
### 3. all-MiniLM-L6-v2
|
||
|
|
|
||
|
|
**Characteristics**:
|
||
|
|
- **Dimensions**: 384 (padded to 768 for schema compatibility)
|
||
|
|
- **Speed**: Very fast (lightweight)
|
||
|
|
- **Quality**: Good for similarity
|
||
|
|
- **Languages**: English
|
||
|
|
- **Use case**: High-throughput, English-only scenarios
|
||
|
|
|
||
|
|
```bash
|
||
|
|
export EMBEDDINGS_MODEL=all-MiniLM-L6-v2
|
||
|
|
# or
|
||
|
|
export EMBEDDINGS_MODEL=sentence-transformers/all-MiniLM-L6-v2
|
||
|
|
```
|
||
|
|
|
||
|
|
### 4. BAAI/bge-small-en-v1.5
|
||
|
|
|
||
|
|
**Characteristics**:
|
||
|
|
- **Dimensions**: 384 (padded to 768)
|
||
|
|
- **Speed**: Very fast
|
||
|
|
- **Quality**: Good for English retrieval
|
||
|
|
- **Languages**: English
|
||
|
|
- **Use case**: Fast English-only systems
|
||
|
|
|
||
|
|
```bash
|
||
|
|
export EMBEDDINGS_MODEL=BAAI/bge-small-en-v1.5
|
||
|
|
```
|
||
|
|
|
||
|
|
### 5. BAAI/bge-base-en-v1.5
|
||
|
|
|
||
|
|
**Characteristics**:
|
||
|
|
- **Dimensions**: 768 (native)
|
||
|
|
- **Speed**: Medium
|
||
|
|
- **Quality**: Excellent for English
|
||
|
|
- **Languages**: English
|
||
|
|
- **Use case**: Best quality for English-only deployments
|
||
|
|
|
||
|
|
```bash
|
||
|
|
export EMBEDDINGS_MODEL=BAAI/bge-base-en-v1.5
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Configuration
|
||
|
|
|
||
|
|
### Environment Variables
|
||
|
|
|
||
|
|
```bash
|
||
|
|
# Select embedding model (default: nomic-ai/nomic-embed-text-v2-moe)
|
||
|
|
export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v2-moe
|
||
|
|
|
||
|
|
# Gateway endpoint (default: https://api.riotpiao.com)
|
||
|
|
export LLM_API_BASE=https://api.riotpiao.com
|
||
|
|
|
||
|
|
# Optional: API key for gated models
|
||
|
|
export LLM_API_KEY=your-api-key
|
||
|
|
```
|
||
|
|
|
||
|
|
### Kubernetes Deployment
|
||
|
|
|
||
|
|
```yaml
|
||
|
|
apiVersion: apps/v1
|
||
|
|
kind: Deployment
|
||
|
|
metadata:
|
||
|
|
name: poimen-memory
|
||
|
|
spec:
|
||
|
|
template:
|
||
|
|
spec:
|
||
|
|
containers:
|
||
|
|
- name: memory
|
||
|
|
env:
|
||
|
|
- name: EMBEDDINGS_MODEL
|
||
|
|
value: "nomic-ai/nomic-embed-text-v1.5" # Higher quality
|
||
|
|
- name: LLM_API_BASE
|
||
|
|
value: "https://api.riotpiao.com"
|
||
|
|
```
|
||
|
|
|
||
|
|
### Docker
|
||
|
|
|
||
|
|
```bash
|
||
|
|
docker run \
|
||
|
|
-e EMBEDDINGS_MODEL=all-MiniLM-L6-v2 \
|
||
|
|
-e LLM_API_BASE=https://api.riotpiao.com \
|
||
|
|
poimen-memory:latest
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Comparison Table
|
||
|
|
|
||
|
|
| Model | Dims | Speed | Quality | Languages | Use Case |
|
||
|
|
|-------|------|-------|---------|-----------|----------|
|
||
|
|
| **nomic-v2-moe** (default) | 768 | ⚡⚡⚡ | ⭐⭐⭐ | 30+ | **Production default** |
|
||
|
|
| **nomic-v1.5** | 768 | ⚡⚡ | ⭐⭐⭐⭐ | 30+ | Quality-first, multilingual |
|
||
|
|
| **all-MiniLM-L6** | 384→768 | ⚡⚡⚡⚡ | ⭐⭐⭐ | English | High throughput |
|
||
|
|
| **bge-small-en** | 384→768 | ⚡⚡⚡⚡ | ⭐⭐⭐ | English | Fast retrieval |
|
||
|
|
| **bge-base-en** | 768 | ⚡⚡ | ⭐⭐⭐⭐⭐ | English | English-only best quality |
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Performance Impact
|
||
|
|
|
||
|
|
### Embedding Latency (per text)
|
||
|
|
|
||
|
|
```
|
||
|
|
nomic-v2-moe: ~50ms ← Default (good balance)
|
||
|
|
all-MiniLM-L6: ~30ms ← Fastest
|
||
|
|
nomic-v1.5: ~80ms
|
||
|
|
bge-base-en: ~60ms
|
||
|
|
```
|
||
|
|
|
||
|
|
### Throughput at 10 batch size
|
||
|
|
|
||
|
|
```
|
||
|
|
nomic-v2-moe: ~200 texts/sec
|
||
|
|
all-MiniLM-L6: ~330 texts/sec ← Highest throughput
|
||
|
|
nomic-v1.5: ~125 texts/sec
|
||
|
|
bge-base-en: ~165 texts/sec
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Switching Models
|
||
|
|
|
||
|
|
### 1. While Running
|
||
|
|
|
||
|
|
```bash
|
||
|
|
# Change env var
|
||
|
|
kubectl set env deployment/poimen-memory \
|
||
|
|
EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v1.5
|
||
|
|
|
||
|
|
# Restart pods (will use new model)
|
||
|
|
kubectl rollout restart deployment/poimen-memory
|
||
|
|
|
||
|
|
# Monitor logs
|
||
|
|
kubectl logs -f deployment/poimen-memory | grep "Embeddings client"
|
||
|
|
# Expected: "Embeddings client initialized: model=nomic-ai/nomic-embed-text-v1.5"
|
||
|
|
```
|
||
|
|
|
||
|
|
### 2. New Ingests
|
||
|
|
|
||
|
|
Changing models only affects **new** ingests. Existing chunks keep their old embeddings.
|
||
|
|
|
||
|
|
To re-embed existing chunks:
|
||
|
|
```bash
|
||
|
|
# 1. Mark all chunks as pending re-embedding
|
||
|
|
psql -h memory-db -U app memory -c \
|
||
|
|
"UPDATE chunks SET indexed_in_pgvector = false, opensearch_pending = true;"
|
||
|
|
|
||
|
|
# 2. Restart queue worker (reprocesses all chunks)
|
||
|
|
kubectl rollout restart deployment/poimen-memory
|
||
|
|
|
||
|
|
# Wait for queue to clear
|
||
|
|
# Monitor: kubectl logs -f deployment/poimen-memory | grep "messages_processed"
|
||
|
|
```
|
||
|
|
|
||
|
|
### 3. Validate Model
|
||
|
|
|
||
|
|
```bash
|
||
|
|
# Check logs for model initialization
|
||
|
|
kubectl logs -n poimen deployment/poimen-memory | grep "Embeddings client"
|
||
|
|
|
||
|
|
# Test embedding endpoint
|
||
|
|
curl -X POST https://api.riotpiao.com/v1/embeddings \
|
||
|
|
-H "Content-Type: application/json" \
|
||
|
|
-d '{
|
||
|
|
"model": "nomic-ai/nomic-embed-text-v1.5",
|
||
|
|
"input": ["hello world"]
|
||
|
|
}' | jq '.data[0].embedding | length'
|
||
|
|
# Output: 768
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Troubleshooting
|
||
|
|
|
||
|
|
### "Unsupported embedding model" Error
|
||
|
|
|
||
|
|
**Error**:
|
||
|
|
```
|
||
|
|
thread 'actix-web' panicked at 'Unsupported embedding model: bert-base-uncased'
|
||
|
|
```
|
||
|
|
|
||
|
|
**Fix**:
|
||
|
|
1. Check supported models list above
|
||
|
|
2. Use one of the validated models
|
||
|
|
3. If you have a custom model, update `validate_model()` in embeddings.rs
|
||
|
|
|
||
|
|
### Wrong Dimensionality
|
||
|
|
|
||
|
|
**Error**:
|
||
|
|
```
|
||
|
|
model custom-embed-384 returned 384-dim vector, expected 768
|
||
|
|
```
|
||
|
|
|
||
|
|
**Cause**: Model returns 384-dim vectors, but pgvector schema expects 768
|
||
|
|
|
||
|
|
**Solutions**:
|
||
|
|
1. Use a model that returns 768-dim (e.g., `nomic-ai/nomic-embed-text-v1.5`)
|
||
|
|
2. Or manually pad 384-dim vectors to 768-dim by adding zeros
|
||
|
|
3. Or re-migrate schema to 384-dim (complex, not recommended)
|
||
|
|
|
||
|
|
### Slow Embedding Performance
|
||
|
|
|
||
|
|
**If latency > 200ms per text**:
|
||
|
|
|
||
|
|
```bash
|
||
|
|
# Try faster model
|
||
|
|
kubectl set env deployment/poimen-memory \
|
||
|
|
EMBEDDINGS_MODEL=all-MiniLM-L6-v2
|
||
|
|
|
||
|
|
# Check embedding service health
|
||
|
|
curl https://api.riotpiao.com/v1/models | jq '.data[] | select(.id | contains("embed"))'
|
||
|
|
|
||
|
|
# Monitor queue worker metrics
|
||
|
|
kubectl logs -f deployment/poimen-memory | grep "processing_time"
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Integration with Queue Worker
|
||
|
|
|
||
|
|
The Queue Worker automatically uses the configured embedding model:
|
||
|
|
|
||
|
|
```rust
|
||
|
|
// crates/mem-cli/src/queue_worker.rs
|
||
|
|
let embedding_vec = embeddings.embed_one(&content).await?;
|
||
|
|
// ↑ Uses model from EMBEDDINGS_MODEL env var
|
||
|
|
```
|
||
|
|
|
||
|
|
When queue worker logs show:
|
||
|
|
```
|
||
|
|
Embeddings client initialized: model=nomic-ai/nomic-embed-text-v1.5
|
||
|
|
```
|
||
|
|
|
||
|
|
All embeddings are computed with that model.
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Validation at Startup
|
||
|
|
|
||
|
|
The system validates model compatibility on HTTP server startup:
|
||
|
|
|
||
|
|
```
|
||
|
|
INFO Embeddings client initialized: model=nomic-ai/nomic-embed-text-v2-moe, base_url=https://api.riotpiao.com
|
||
|
|
```
|
||
|
|
|
||
|
|
If validation fails, server refuses to start:
|
||
|
|
|
||
|
|
```
|
||
|
|
ERROR Unsupported embedding model: unknown-model. Supported models: [...]
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Adding Custom Models
|
||
|
|
|
||
|
|
To support a new embedding model:
|
||
|
|
|
||
|
|
1. **Verify dimension**: Test with gateway
|
||
|
|
```bash
|
||
|
|
curl -X POST https://api.riotpiao.com/v1/embeddings \
|
||
|
|
-d '{"model": "my-custom-model", "input": ["test"]}'
|
||
|
|
# Check dimension count in response
|
||
|
|
```
|
||
|
|
|
||
|
|
2. **Add to allowed list** (crates/mem-llm/src/embeddings.rs):
|
||
|
|
```rust
|
||
|
|
let supported_models = vec![
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
"my-custom-model", // ← Add here
|
||
|
|
];
|
||
|
|
```
|
||
|
|
|
||
|
|
3. **Document dimensions** in this file
|
||
|
|
|
||
|
|
4. **Test**:
|
||
|
|
```bash
|
||
|
|
EMBEDDINGS_MODEL=my-custom-model cargo run --bin mem -- serve
|
||
|
|
```
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## References
|
||
|
|
|
||
|
|
- [Nomic AI Models](https://www.nomic.ai/)
|
||
|
|
- [BGE Models (BAAI)](https://huggingface.co/BAAI/bge-base-en-v1.5)
|
||
|
|
- [Sentence Transformers](https://www.sbert.net/)
|
||
|
|
- [Gateway Embeddings API](https://api.riotpiao.com/docs#/embeddings)
|