Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2bccb9680d | |||
| 0842310eba | |||
| 95d458ce06 | |||
| 5345ed5e0b |
@@ -61,11 +61,13 @@ All rerankers share these common configuration parameters:
|
||||
|
||||
### LLM Reranker
|
||||
|
||||
| Parameter | Description | Type | Default |
|
||||
| -------------- | --------------------------- | ------ | -------- |
|
||||
| `llm.provider` | LLM provider for reranking | `str` | Required |
|
||||
| `llm.config` | LLM configuration object | `dict` | Required |
|
||||
| `top_n` | Number of results to return | `int` | `None` |
|
||||
| Parameter | Description | Type | Default |
|
||||
| ----------- | -------------------------------------------------------------------------------------------- | ------ | -------------- |
|
||||
| `provider` | LLM provider (`openai`, `anthropic`, etc.) | `str` | `"openai"` |
|
||||
| `model` | LLM model to use for scoring | `str` | `"gpt-4o-mini"` |
|
||||
| `api_key` | API key for LLM provider | `str` | `None` |
|
||||
| `top_k` | Number of results to return after reranking | `int` | `None` |
|
||||
| `llm` | Optional nested LLM config (`{"provider": ..., "config": {...}}`) for provider-specific fields (e.g. `ollama_base_url`) | `dict` | `None` |
|
||||
|
||||
## Environment Variables
|
||||
|
||||
|
||||
@@ -27,6 +27,10 @@ Provide only a single numerical score between 0.0 and 1.0. Do not include any ex
|
||||
|
||||
## Custom Prompt Configuration
|
||||
|
||||
<Warning>
|
||||
The `scoring_prompt` field is **deprecated** and will be removed in a future version. Setting it still works but emits a `DeprecationWarning` at runtime. The field currently sets the system prompt used by the LLM; custom prompts may be exposed through a different API in a future release.
|
||||
</Warning>
|
||||
|
||||
You can provide a custom prompt template using the `scoring_prompt` parameter:
|
||||
|
||||
```python
|
||||
@@ -189,7 +193,7 @@ for i, prompt in enumerate(prompts):
|
||||
config["reranker"]["config"]["scoring_prompt"] = prompt
|
||||
memory = Memory.from_config(config)
|
||||
|
||||
results = memory.search("test query", filters={"user_id": "test_user"})
|
||||
results = memory.search("test query", filters={"user_id": "test_user"}, rerank=True)
|
||||
print(f"Prompt {i+1} results: {results}")
|
||||
```
|
||||
|
||||
|
||||
@@ -74,7 +74,7 @@ os.environ["COHERE_API_KEY"] = "your-api-key"
|
||||
config = {
|
||||
"vector_store": {"provider": "chroma"},
|
||||
"llm": {"provider": "openai", "config": {"model": "gpt-4o-mini"}},
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "cohere",
|
||||
"config": {
|
||||
"model": "rerank-english-v3.0",
|
||||
@@ -95,7 +95,7 @@ messages = [
|
||||
memory.add(messages, user_id="bob")
|
||||
|
||||
# Search with reranking
|
||||
results = memory.search("What is the user's profession?", filters={"user_id": "bob"})
|
||||
results = memory.search("What is the user's profession?", filters={"user_id": "bob"}, rerank=True)
|
||||
|
||||
for result in results['results']:
|
||||
print(f"Memory: {result['memory']}")
|
||||
@@ -110,7 +110,7 @@ For multilingual applications, use the multilingual model:
|
||||
|
||||
```python Python
|
||||
config = {
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "cohere",
|
||||
"config": {
|
||||
"model": "rerank-multilingual-v3.0",
|
||||
|
||||
@@ -35,7 +35,7 @@ m = Memory.from_config(config)
|
||||
| `device` | str | "cpu" | Device to run model on ("cpu", "cuda", "mps") |
|
||||
| `batch_size` | int | 32 | Batch size for processing |
|
||||
| `max_length` | int | 512 | Maximum input sequence length |
|
||||
| `trust_remote_code` | bool | False | Allow remote code execution |
|
||||
| `normalize` | bool | True | Whether to normalize scores to [0, 1] range |
|
||||
|
||||
### Advanced Configuration
|
||||
|
||||
@@ -48,10 +48,7 @@ config = {
|
||||
"device": "cuda",
|
||||
"batch_size": 16,
|
||||
"max_length": 512,
|
||||
"trust_remote_code": False,
|
||||
"model_kwargs": {
|
||||
"torch_dtype": "float16"
|
||||
}
|
||||
"normalize": True
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -154,7 +151,7 @@ m.add("I enjoy reading science fiction books", user_id="alice")
|
||||
# Search with reranking
|
||||
results = m.search(
|
||||
"What outdoor activities do I enjoy?",
|
||||
user_id="alice",
|
||||
filters={"user_id": "alice"},
|
||||
rerank=True
|
||||
)
|
||||
|
||||
@@ -208,10 +205,7 @@ config = {
|
||||
"model": "BAAI/bge-reranker-base",
|
||||
"device": "cpu",
|
||||
"batch_size": 8, # Smaller batch size
|
||||
"max_length": 256, # Shorter sequences
|
||||
"model_kwargs": {
|
||||
"torch_dtype": "float16" # Half precision
|
||||
}
|
||||
"max_length": 256 # Shorter sequences
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -232,7 +226,7 @@ config = {
|
||||
try:
|
||||
results = m.search(
|
||||
"test query",
|
||||
user_id="alice",
|
||||
filters={"user_id": "alice"},
|
||||
rerank=True
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -240,7 +234,7 @@ except Exception as e:
|
||||
# Fall back to vector search only
|
||||
results = m.search(
|
||||
"test query",
|
||||
user_id="alice",
|
||||
filters={"user_id": "alice"},
|
||||
rerank=False
|
||||
)
|
||||
```
|
||||
@@ -251,13 +245,13 @@ except Exception as e:
|
||||
|
||||
```python
|
||||
# Use a private model from Hugging Face
|
||||
# Set HF_TOKEN environment variable for private model access
|
||||
config = {
|
||||
"reranker": {
|
||||
"provider": "huggingface",
|
||||
"config": {
|
||||
"model": "your-org/custom-reranker",
|
||||
"device": "cuda",
|
||||
"use_auth_token": "your-hf-token"
|
||||
"device": "cuda"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -311,13 +305,13 @@ config = {
|
||||
import os
|
||||
os.environ["TRANSFORMERS_CACHE"] = "/path/to/cache"
|
||||
|
||||
# Or use offline mode
|
||||
# Or use offline mode by setting the environment variable
|
||||
# TRANSFORMERS_OFFLINE=1 before running your script
|
||||
config = {
|
||||
"reranker": {
|
||||
"provider": "huggingface",
|
||||
"config": {
|
||||
"model": "BAAI/bge-reranker-base",
|
||||
"local_files_only": True
|
||||
"model": "BAAI/bge-reranker-base"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -268,7 +268,7 @@ m.add("I'm vegetarian", user_id="alice")
|
||||
# Search with LLM reranking
|
||||
results = m.search(
|
||||
"What foods should I avoid?",
|
||||
user_id="alice",
|
||||
filters={"user_id": "alice"},
|
||||
rerank=True
|
||||
)
|
||||
|
||||
|
||||
@@ -39,7 +39,7 @@ config = {
|
||||
"model": "gpt-4o-mini"
|
||||
}
|
||||
},
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
@@ -60,7 +60,7 @@ For better performance, use GPU acceleration:
|
||||
|
||||
```python Python
|
||||
config = {
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
@@ -80,7 +80,7 @@ from mem0 import Memory
|
||||
config = {
|
||||
"vector_store": {"provider": "chroma"},
|
||||
"llm": {"provider": "openai", "config": {"model": "gpt-4o-mini"}},
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
@@ -101,7 +101,7 @@ messages = [
|
||||
memory.add(messages, user_id="charlie")
|
||||
|
||||
# Search with local reranking
|
||||
results = memory.search("What books does the user like?", filters={"user_id": "charlie"})
|
||||
results = memory.search("What books does the user like?", filters={"user_id": "charlie"}, rerank=True)
|
||||
|
||||
for result in results['results']:
|
||||
print(f"Memory: {result['memory']}")
|
||||
@@ -117,7 +117,7 @@ You can use any HuggingFace cross-encoder model:
|
||||
```python Python
|
||||
# Using a different model
|
||||
config = {
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"model": "cross-encoder/stsb-distilroberta-base",
|
||||
|
||||
@@ -37,7 +37,7 @@ config = {
|
||||
"model": "gpt-4o-mini"
|
||||
}
|
||||
},
|
||||
"rerank": {
|
||||
"reranker": {
|
||||
"provider": "zero_entropy",
|
||||
"config": {
|
||||
"model": "zerank-1", # or "zerank-1-small"
|
||||
@@ -71,7 +71,7 @@ os.environ["ZERO_ENTROPY_API_KEY"] = "your-api-key"
|
||||
config = {
|
||||
"vector_store": {"provider": "chroma"},
|
||||
"llm": {"provider": "openai", "config": {"model": "gpt-4o-mini"}},
|
||||
"rerank": {"provider": "zero_entropy", "config": {"model": "zerank-1"}}
|
||||
"reranker": {"provider": "zero_entropy", "config": {"model": "zerank-1"}}
|
||||
}
|
||||
|
||||
memory = Memory.from_config(config)
|
||||
@@ -86,7 +86,7 @@ messages = [
|
||||
memory.add(messages, user_id="alice")
|
||||
|
||||
# Search with reranking
|
||||
results = memory.search("What Italian food does the user like?", filters={"user_id": "alice"})
|
||||
results = memory.search("What Italian food does the user like?", filters={"user_id": "alice"}, rerank=True)
|
||||
|
||||
for result in results['results']:
|
||||
print(f"Memory: {result['memory']}")
|
||||
|
||||
@@ -13,10 +13,10 @@ The number of candidates sent to the reranker significantly impacts performance:
|
||||
```python
|
||||
# Optimal candidate sizes for different rerankers
|
||||
config_map = {
|
||||
"cohere": {"initial_candidates": 100, "top_n": 10},
|
||||
"sentence_transformer": {"initial_candidates": 50, "top_n": 10},
|
||||
"huggingface": {"initial_candidates": 30, "top_n": 5},
|
||||
"llm_reranker": {"initial_candidates": 20, "top_n": 5}
|
||||
"cohere": {"initial_candidates": 100, "top_k": 10},
|
||||
"sentence_transformer": {"initial_candidates": 50, "top_k": 10},
|
||||
"huggingface": {"initial_candidates": 30, "top_k": 5},
|
||||
"llm_reranker": {"initial_candidates": 20, "top_k": 5}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -31,7 +31,7 @@ config = {
|
||||
"config": {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
"batch_size": 16, # Process multiple candidates at once
|
||||
"top_n": 10
|
||||
"top_k": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -48,7 +48,7 @@ config = {
|
||||
"provider": "cohere",
|
||||
"config": {
|
||||
"model": "rerank-english-v3.0",
|
||||
"top_n": 10,
|
||||
"top_k": 10,
|
||||
"max_chunks_per_doc": 10, # Limit chunk processing
|
||||
"return_documents": False # Reduce response size
|
||||
}
|
||||
@@ -73,7 +73,7 @@ config = {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
"device": "cuda", # Use GPU when available
|
||||
"batch_size": 32,
|
||||
"top_n": 10,
|
||||
"top_k": 10,
|
||||
"max_length": 512 # Limit input length
|
||||
}
|
||||
}
|
||||
@@ -107,10 +107,9 @@ config = {
|
||||
"provider": "huggingface",
|
||||
"config": {
|
||||
"model": "BAAI/bge-reranker-base",
|
||||
"use_fp16": True, # Half precision for speed
|
||||
"max_length": 512,
|
||||
"batch_size": 8,
|
||||
"top_n": 10
|
||||
"top_k": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -124,17 +123,11 @@ config = {
|
||||
"reranker": {
|
||||
"provider": "llm_reranker",
|
||||
"config": {
|
||||
"llm": {
|
||||
"provider": "openai",
|
||||
"config": {
|
||||
"model": "gpt-3.5-turbo", # Faster than gpt-4
|
||||
"temperature": 0, # Deterministic results
|
||||
"max_tokens": 500 # Limit response length
|
||||
}
|
||||
},
|
||||
"batch_ranking": True, # Rank multiple at once
|
||||
"top_n": 5, # Fewer results for faster processing
|
||||
"timeout": 10 # Request timeout
|
||||
"provider": "openai",
|
||||
"model": "gpt-4o-mini", # Faster than gpt-4o
|
||||
"temperature": 0.0, # Deterministic results
|
||||
"max_tokens": 100, # Limit response length
|
||||
"top_k": 5 # Fewer results for faster processing
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -153,7 +146,7 @@ def measure_reranker_performance(config, queries, user_id):
|
||||
latencies = []
|
||||
for query in queries:
|
||||
start_time = time.time()
|
||||
results = memory.search(query, filters={"user_id": user_id})
|
||||
results = memory.search(query, filters={"user_id": user_id}, rerank=True)
|
||||
latency = time.time() - start_time
|
||||
latencies.append(latency)
|
||||
|
||||
@@ -201,12 +194,12 @@ class CachedReranker:
|
||||
### Model Caching
|
||||
```python
|
||||
# Pre-load models to avoid initialization overhead
|
||||
# Set SENTENCE_TRANSFORMERS_HOME env var to control cache location
|
||||
config = {
|
||||
"reranker": {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
"cache_folder": "/path/to/model/cache",
|
||||
"device": "cuda"
|
||||
}
|
||||
}
|
||||
@@ -218,14 +211,14 @@ config = {
|
||||
### Async Configuration
|
||||
```python
|
||||
import asyncio
|
||||
from mem0 import Memory
|
||||
from mem0 import AsyncMemory
|
||||
|
||||
async def parallel_search(config, queries, user_id):
|
||||
memory = Memory.from_config(config)
|
||||
memory = AsyncMemory.from_config(config)
|
||||
|
||||
# Process multiple queries concurrently
|
||||
tasks = [
|
||||
memory.search_async(query, user_id=user_id)
|
||||
memory.search(query, filters={"user_id": user_id}, rerank=True)
|
||||
for query in queries
|
||||
]
|
||||
|
||||
@@ -249,8 +242,7 @@ config = {
|
||||
"config": {
|
||||
"device": "cuda",
|
||||
"model": "cross-encoder/ms-marco-electra-base",
|
||||
"batch_size": 64, # Larger batch for GPU
|
||||
"fp16": True # Half precision
|
||||
"batch_size": 64 # Larger batch for GPU
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -268,8 +260,7 @@ config = {
|
||||
"provider": "sentence_transformer",
|
||||
"config": {
|
||||
"device": "cpu",
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
"num_workers": 4 # Parallel processing
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user