README.md

---

license: apache-2.0
language:
- en
tags:
- sentence-transformers
- embeddings
- retrieval
- agents
- memory
- rag
- semantic-search
- ai-agents
- llm-memory
- vector-search
library_name: transformers
pipeline_tag: sentence-similarity
datasets:
- custom
metrics:
- mrr
- recall
- ndcg
model-index:
- name: agentrank-base
  results:
  - task:
      type: retrieval
      name: Agent Memory Retrieval
    metrics:
    - type: mrr
      value: 0.6496
      name: MRR
    - type: recall
      value: 0.4440
      name: Recall@1
    - type: recall
      value: 0.9960
      name: Recall@5
    - type: ndcg
      value: 0.6786
      name: NDCG@10
---


<div align="center">

# 🧠 AgentRank-Base

### The First Embedding Model Built Specifically for AI Agent Memory Retrieval

<p>
  <img src="https://img.shields.io/badge/MRR-0.65-brightgreen?style=for-the-badge" alt="MRR">
  <img src="https://img.shields.io/badge/Recall%405-99.6%25-blue?style=for-the-badge" alt="Recall@5">
  <img src="https://img.shields.io/badge/Parameters-149M-orange?style=for-the-badge" alt="Parameters">
  <img src="https://img.shields.io/badge/License-Apache%202.0-green?style=for-the-badge" alt="License">
</p>

**+23% MRR improvement over general-purpose embedders** | **Temporal awareness** | **Memory type understanding**

[🚀 Quick Start](#-quick-start) • [📊 Benchmarks](#-benchmarks) • [🔧 Architecture](#-architecture) • [💡 Why AgentRank?](#-why-agentrank)

</div>

---

## 🎯 TL;DR

> **AgentRank-Base** is an embedding model designed for AI agents that need to remember. Unlike generic embedders (OpenAI, Cohere, MiniLM), AgentRank understands:
> - ⏰ **When** something happened (temporal awareness)
> - 📁 **What type** of memory it is (episodic vs semantic vs procedural)
> - ⭐ **How important** the memory is

---

## 💡 Why AgentRank?

### The Problem with Current Embedders

AI agents need memory. But when you ask an agent:

> *"What did we discuss about Python **yesterday**?"*

Current embedders fail because they:
- ❌ Don't understand "yesterday" means recent time
- ❌ Can't distinguish between a preference and an event
- ❌ Treat all memories as equally important

### The AgentRank Solution

| Challenge | OpenAI/Cohere/MiniLM | AgentRank |
|-----------|---------------------|-----------|
| "What did I say **yesterday**?" | Random old results 😕 | Recent memories first ✅ |
| "What's my **preference**?" | Mixed with events 😕 | Only preferences ✅ |
| "What's **most important**?" | No priority 😕 | Importance-aware retrieval ✅ |

---

## 📊 Benchmarks

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

| Model | Parameters | MRR ↑ | Recall@1 ↑ | Recall@5 ↑ | NDCG@10 ↑ |
|-------|------------|-------|------------|------------|-----------|
| **AgentRank-Base** | 149M | **0.6496** | **0.4440** | **0.9960** | **0.6786** |
| AgentRank-Small | 33M | 0.6375 | 0.4460 | 0.9740 | 0.6797 |
| all-mpnet-base-v2 | 109M | 0.5351 | 0.3660 | 0.7960 | 0.6335 |
| all-MiniLM-L6-v2 | 22M | 0.5297 | 0.3720 | 0.7520 | 0.6370 |

### Improvement Over Baselines

| vs Baseline | MRR | Recall@1 | Recall@5 |
|-------------|-----|----------|----------|
| vs MiniLM | **+22.6%** | **+19.4%** | **+32.4%** |
| vs MPNet | **+21.4%** | **+21.3%** | **+25.1%** |

---

## 🚀 Quick Start

### Installation

```bash

pip install transformers torch

```

### Basic Usage

```python

from transformers import AutoModel, AutoTokenizer

import torch



# Load model and tokenizer

model = AutoModel.from_pretrained("vrushket/agentrank-base")

tokenizer = AutoTokenizer.from_pretrained("vrushket/agentrank-base")



def encode(texts, model, tokenizer):

    """Encode texts to embeddings."""

    inputs = tokenizer(

        texts, 

        padding=True, 

        truncation=True, 

        max_length=512,

        return_tensors="pt"

    )

    with torch.no_grad():

        outputs = model(**inputs)

        # Mean pooling

        embeddings = outputs.last_hidden_state.mean(dim=1)

        # L2 normalize

        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

    return embeddings



# Your agent's memories

memories = [

    "User prefers Python over JavaScript for backend development",

    "User asked about React frameworks yesterday",

    "User mentioned they have 3 years of coding experience",

    "User is working on an e-commerce project",

]



# A query from the user

query = "What programming language does the user prefer?"



# Encode everything

memory_embeddings = encode(memories, model, tokenizer)

query_embedding = encode([query], model, tokenizer)



# Find most similar memory

similarities = torch.mm(query_embedding, memory_embeddings.T)[0]

best_match_idx = similarities.argmax().item()



print(f"Query: {query}")

print(f"Best match: {memories[best_match_idx]}")

print(f"Similarity: {similarities[best_match_idx]:.4f}")



# Output:

# Query: What programming language does the user prefer?

# Best match: User prefers Python over JavaScript for backend development

# Similarity: 0.8234

```

### Advanced Usage with Metadata

For full temporal and memory type awareness, use the AgentRank package:

```python

# Coming soon: pip install agentrank

from agentrank import AgentRankEmbedder



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



# Encode with temporal context

memory_embedding = model.encode(

    text="User mentioned they prefer morning meetings",

    days_ago=7,           # Memory is 1 week old

    memory_type="semantic" # It's a preference (not an event)

)



# Encode query (no metadata needed for queries)

query_embedding = model.encode("When does the user like to have meetings?")



# The model now knows this is a week-old preference!

similarity = torch.cosine_similarity(query_embedding, memory_embedding, dim=0)

```

---

## 🔧 Architecture

AgentRank-Base is built on **ModernBERT-base** (110M params) with novel additions:

```

┌─────────────────────────────────────────────────┐

│     ModernBERT Encoder (22 Transformer Layers)  │

│     - RoPE Positional Encoding                  │

│     - Flash Attention                           │

│     - 768 Hidden Dimension                      │

└─────────────────────────────────────────────────┘

                       │

       ┌───────────────┼───────────────┐

       ↓               ↓               ↓

┌─────────────┐ ┌─────────────┐ ┌─────────────┐

│  Temporal   │ │  Memory     │ │ Importance  │

│  Position   │ │  Type       │ │ Prediction  │

│  Embeddings │ │  Embeddings │ │ Head        │

│  (10 × 768) │ │  (4 × 768)  │ │ (768→1)     │

└─────────────┘ └─────────────┘ └─────────────┘

       │               │               │

       └───────────────┼───────────────┘

                       ↓

          ┌─────────────────────┐

          │  Projection Layer   │

          │  (768 → 768)        │

          └─────────────────────┘

                       ↓

          ┌─────────────────────┐

          │  L2 Normalization   │

          │  768-dim Embedding  │

          └─────────────────────┘

```

### Novel Components

| Component | Purpose | How It Helps |
|-----------|---------|--------------|
| **Temporal Embeddings** | Encodes memory age (today, this week, last month, etc.) | "Yesterday" queries match recent memories |
| **Memory Type Embeddings** | Distinguishes episodic/semantic/procedural | "What do I like?" matches preferences, not events |
| **Importance Head** | Auxiliary task predicting memory priority | Helps learn better representations |

### Temporal Buckets

```

Bucket 0: Today (0-1 days)

Bucket 1: Recent (1-3 days)

Bucket 2: This week (3-7 days)

Bucket 3: Last week (7-14 days)

Bucket 4: This month (14-30 days)

Bucket 5: Last month (30-60 days)

Bucket 6: Few months (60-90 days)

Bucket 7: Half year (90-180 days)

Bucket 8: This year (180-365 days)

Bucket 9: Long ago (365+ days)

```

### Memory Types

```

Type 0: Episodic   → Events, conversations ("We discussed X yesterday")

Type 1: Semantic   → Facts, preferences ("User likes Python")

Type 2: Procedural → Instructions ("To deploy, run npm build")

Type 3: Unknown    → Fallback

```

---

## 🎓 Training Details

| Aspect | Details |
|--------|---------|
| **Base Model** | answerdotai/ModernBERT-base (110M params) |
| **Training Data** | 500K synthetic agent memory samples |
| **Memory Distribution** | Episodic (40%), Semantic (35%), Procedural (25%) |
| **Loss Function** | Multiple Negatives Ranking Loss + Importance MSE |
| **Hard Negatives** | 7 per sample (5 types: temporal, type confusion, topic drift, etc.) |
| **Batch Size** | 16-32 per GPU |
| **Hardware** | 2× NVIDIA RTX 6000 Ada (48GB each) |
| **Training Time** | ~12 hours |
| **Precision** | FP16 Mixed Precision |
| **Final Val Loss** | 0.877 |

---

## 🏗️ Use Cases

### 1. AI Agents with Long-Term Memory

```python

# Store memories with metadata

agent.remember(

    text="User is allergic to peanuts",

    memory_type="semantic",

    importance=10,  # Critical medical info!

)



# Later, when discussing food...

relevant_memories = agent.recall("What should I know about the user's diet?")

# Returns: "User is allergic to peanuts" (even if stored months ago)

```

### 2. RAG Systems for Conversational AI

```python

# Better retrieval for chatbots

query = "What did we talk about in our last meeting?"

# AgentRank returns recent, relevant conversations

# Generic embedders return random topically-similar docs

```

### 3. Personal Knowledge Bases

```python

# User's notes and preferences

memories = [

    "I prefer dark mode in all apps",

    "My morning routine starts at 6 AM",

    "Important: Tax deadline April 15",

]

# AgentRank properly handles time-sensitive queries

```

---

## 🆚 When to Use AgentRank vs Others

| Use Case | Best Model |
|----------|------------|
| **AI agents with memory** | ✅ AgentRank |
| **Time-sensitive retrieval** | ✅ AgentRank |
| **Conversational AI** | ✅ AgentRank |
| General document search | OpenAI / Cohere |
| Code search | CodeBERT |
| Scientific papers | SciBERT |

---

## 📁 Model Family

| Model | Parameters | Speed | Quality | Best For |
|-------|------------|-------|---------|----------|
| [agentrank-small](https://huggingface.co/vrushket/agentrank-small) | 33M | ⚡⚡⚡ Fast | Good | Real-time agents, edge |
| **agentrank-base** | 149M | ⚡⚡ Medium | **Best** | Quality-critical apps |
| agentrank-reranker (coming) | 149M | ⚡ Slower | Superior | Two-stage retrieval |

---

## 📚 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-base}

}

```

---

## 🤝 Community & Support

- 🐛 **Issues**: [GitHub Issues](https://github.com/vmore2/AgentRank-base/issues)
- 💬 **Discussions**: [HuggingFace Community](https://huggingface.co/vrushket/agentrank-base/discussions)
- 📧 **Contact**: vrushket2604@gmail.com

---

## 📄 License

Apache 2.0 - **Free for commercial use!**

---

<div align="center">

### ⭐ If AgentRank helps your project, please star the repo!

**Built with ❤️ for the AI agent community**

</div>