README.md

---
license: apache-2.0
language:
- en
tags:
- sentence-transformers
- embeddings
- retrieval
- agents
- memory
- rag
- semantic-search
library_name: transformers
pipeline_tag: sentence-similarity
datasets:
- custom
metrics:
- mrr
- recall
- ndcg
model-index:
- name: agentrank-small
  results:
  - task:
      type: retrieval
      name: Agent Memory Retrieval
    metrics:
    - type: mrr
      value: 0.6375
      name: MRR
    - type: recall
      value: 0.4460
      name: Recall@1
    - type: recall
      value: 0.9740
      name: Recall@5
    - type: ndcg
      value: 0.6797
      name: NDCG@10
---

# AgentRank-Small: Embedding Model for AI Agent Memory Retrieval

<p align="center">
  <img src="https://img.shields.io/badge/MRR-0.6375-brightgreen" alt="MRR">
  <img src="https://img.shields.io/badge/Recall%405-97.4%25-blue" alt="Recall@5">
  <img src="https://img.shields.io/badge/Parameters-33M-orange" alt="Parameters">
  <img src="https://img.shields.io/badge/License-Apache%202.0-green" alt="License">
</p>

**AgentRank** is the first embedding model family specifically designed for AI agent memory retrieval. Unlike general-purpose embedders, AgentRank understands temporal context, memory types, and importance - critical for agents that need to remember past interactions.

## 🚀 Key Results

| Model | MRR | Recall@1 | Recall@5 | NDCG@10 |
|-------|-----|----------|----------|---------|
| **AgentRank-Small** | **0.6375** | **0.4460** | **0.9740** | **0.6797** |
| all-MiniLM-L6-v2 | 0.5297 | 0.3720 | 0.7520 | 0.6370 |
| all-mpnet-base-v2 | 0.5351 | 0.3660 | 0.7960 | 0.6335 |

**+20% MRR improvement over base MiniLM model!**

## 🎯 Why AgentRank?

AI agents need memory that understands:

| Challenge | General Embedders | AgentRank |
|-----------|-------------------|-----------|
| "What did I say **yesterday**?" | ❌ No temporal awareness | ✅ Temporal embeddings |
| "What's my **preference**?" | ❌ Mixes with events | ✅ Memory type awareness |
| "What's **most important**?" | ❌ No priority | ✅ Importance prediction |

## 📦 Installation

```bash
pip install transformers torch
```

## 💻 Usage

### Basic Usage

```python
from transformers import AutoModel, AutoTokenizer
import torch

# Load model
model = AutoModel.from_pretrained("vrushket/agentrank-small")
tokenizer = AutoTokenizer.from_pretrained("vrushket/agentrank-small")

def encode(texts):
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
    return embeddings

# Encode memories and query
memories = [
    "User prefers Python over JavaScript",
    "User asked about machine learning yesterday",
    "User is working on a web project",
]
query = "What programming language does the user like?"

memory_embeddings = encode(memories)
query_embedding = encode([query])

# Compute similarities
similarities = torch.mm(query_embedding, memory_embeddings.T)
print(f"Most relevant: {memories[similarities.argmax()]}")
# Output: "User prefers Python over JavaScript"
```

### With Temporal & Memory Type Metadata (Full Power)

```python
# For full AgentRank features including temporal awareness:
# pip install agentrank  (coming soon!)

from agentrank import AgentRankEmbedder

model = AgentRankEmbedder.from_pretrained("vrushket/agentrank-small")

# Encode with metadata
embedding = model.encode(
    "User mentioned they prefer morning meetings",
    days_ago=3,           # Memory is 3 days old
    memory_type="semantic" # It's a preference, not an event
)
```

## 🏗️ Architecture

AgentRank-Small is based on `all-MiniLM-L6-v2` with novel additions:

```
┌─────────────────────────────────────────┐
│  MiniLM Transformer Encoder (6 layers)  │
└─────────────────────────────────────────┘
                    │
    ┌───────────────┼───────────────┐
    ↓               ↓               ↓
┌─────────┐   ┌──────────┐   ┌───────────┐
│ Temporal │   │ Memory   │   │ Importance│
│ Position │   │ Type     │   │ Prediction│
│ Embed    │   │ Embed    │   │ Head      │
└─────────┘   └──────────┘   └───────────┘
    │               │               │
    └───────────────┼───────────────┘
                    ↓
         ┌─────────────────┐
         │ L2 Normalized   │
         │ 384-dim Embedding│
         └─────────────────┘
```

**Novel Features:**
- **Temporal Position Embeddings**: 10 learnable buckets (today, 1-3 days, week, month, etc.)
- **Memory Type Embeddings**: Episodic, Semantic, Procedural
- **Importance Prediction Head**: Auxiliary task during training

## 🎓 Training

- **Dataset**: 500K synthetic agent memory samples
- **Memory Types**: Episodic (40%), Semantic (35%), Procedural (25%)
- **Loss**: Multiple Negatives Ranking Loss + Importance MSE
- **Hard Negatives**: 5 types (temporal, type confusion, topic drift, etc.)
- **Hardware**: NVIDIA RTX 6000 Ada (48GB) with FP16

## 📊 Benchmarks

Evaluated on AgentMemBench (500 test samples, 8 candidates each):

| Metric | AgentRank-Small | MiniLM | Improvement |
|--------|-----------------|--------|-------------|
| MRR | 0.6375 | 0.5297 | **+20.4%** |
| Recall@1 | 0.4460 | 0.3720 | **+19.9%** |
| Recall@5 | 0.9740 | 0.7520 | **+29.5%** |
| NDCG@10 | 0.6797 | 0.6370 | **+6.7%** |

## 🔜 Coming Soon

- **AgentRank-Base**: 110M params, even better performance
- **AgentRank-Reranker**: Cross-encoder for top-k refinement
- **Python Package**: `pip install agentrank`

## 📚 Citation

```bibtex
@misc{agentrank2024,
  author = {Vrushket More},
  title = {AgentRank: Embedding Models for AI Agent Memory Retrieval},
  year = {2024},
  publisher = {HuggingFace},
  url = {https://huggingface.co/vrushket/agentrank-small}
}
```

## 📄 License

Apache 2.0 - Free for commercial use!

## 🤝 Acknowledgments

Built on top of [sentence-transformers](https://www.sbert.net/) and [MiniLM](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2).