vrushket/agentrank-base

Model

1

stars

3

commits

1

linked in READMEs

Dec 22, 2025

updated

agents
ai-agents
embeddings
endpoints_compatible
feature-extraction
llm-memory
memory
model-index
modernbert
rag
retrieval
safetensors
semantic-search
sentence-similarity
sentence-transformers
text-embeddings-inference
transformers
vector-search

README

🧠 AgentRank-Base

The First Embedding Model Built Specifically for AI Agent Memory Retrieval

MRR Recall@5 Parameters License

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

πŸš€ Quick Start β€’ πŸ“Š Benchmarks β€’ πŸ”§ Architecture β€’ πŸ’‘ Why AgentRank?


🎯 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

ChallengeOpenAI/Cohere/MiniLMAgentRank
"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):

ModelParametersMRR ↑Recall@1 ↑Recall@5 ↑NDCG@10 ↑
AgentRank-Base149M0.64960.44400.99600.6786
AgentRank-Small33M0.63750.44600.97400.6797
all-mpnet-base-v2109M0.53510.36600.79600.6335
all-MiniLM-L6-v222M0.52970.37200.75200.6370

Improvement Over Baselines

vs BaselineMRRRecall@1Recall@5
vs MiniLM+22.6%+19.4%+32.4%
vs MPNet+21.4%+21.3%+25.1%

πŸš€ Quick Start

Installation

pip install transformers torch

Basic Usage

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:

# 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

ComponentPurposeHow It Helps
Temporal EmbeddingsEncodes memory age (today, this week, last month, etc.)"Yesterday" queries match recent memories
Memory Type EmbeddingsDistinguishes episodic/semantic/procedural"What do I like?" matches preferences, not events
Importance HeadAuxiliary task predicting memory priorityHelps 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

AspectDetails
Base Modelanswerdotai/ModernBERT-base (110M params)
Training Data500K synthetic agent memory samples
Memory DistributionEpisodic (40%), Semantic (35%), Procedural (25%)
Loss FunctionMultiple Negatives Ranking Loss + Importance MSE
Hard Negatives7 per sample (5 types: temporal, type confusion, topic drift, etc.)
Batch Size16-32 per GPU
Hardware2Γ— NVIDIA RTX 6000 Ada (48GB each)
Training Time~12 hours
PrecisionFP16 Mixed Precision
Final Val Loss0.877

πŸ—οΈ Use Cases

1. AI Agents with Long-Term Memory

# 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

# 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

# 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 CaseBest Model
AI agents with memoryβœ… AgentRank
Time-sensitive retrievalβœ… AgentRank
Conversational AIβœ… AgentRank
General document searchOpenAI / Cohere
Code searchCodeBERT
Scientific papersSciBERT

πŸ“ Model Family

ModelParametersSpeedQualityBest For
agentrank-small33M⚑⚑⚑ FastGoodReal-time agents, edge
agentrank-base149M⚑⚑ MediumBestQuality-critical apps
agentrank-reranker (coming)149M⚑ SlowerSuperiorTwo-stage retrieval

πŸ“š Citation

@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


πŸ“„ License

Apache 2.0 - Free for commercial use!


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

Built with ❀️ for the AI agent community

Contributors

vrushket

3 commits

vrushket/agentrank-base

Model

1

stars

3

commits

1

linked in READMEs

Dec 22, 2025

updated

agents
ai-agents
embeddings
endpoints_compatible
feature-extraction
llm-memory
memory
model-index
modernbert
rag
retrieval
safetensors
semantic-search
sentence-similarity
sentence-transformers
text-embeddings-inference
transformers
vector-search

README

🧠 AgentRank-Base

The First Embedding Model Built Specifically for AI Agent Memory Retrieval

MRR Recall@5 Parameters License

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

πŸš€ Quick Start β€’ πŸ“Š Benchmarks β€’ πŸ”§ Architecture β€’ πŸ’‘ Why AgentRank?


🎯 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

ChallengeOpenAI/Cohere/MiniLMAgentRank
"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):

ModelParametersMRR ↑Recall@1 ↑Recall@5 ↑NDCG@10 ↑
AgentRank-Base149M0.64960.44400.99600.6786
AgentRank-Small33M0.63750.44600.97400.6797
all-mpnet-base-v2109M0.53510.36600.79600.6335
all-MiniLM-L6-v222M0.52970.37200.75200.6370

Improvement Over Baselines

vs BaselineMRRRecall@1Recall@5
vs MiniLM+22.6%+19.4%+32.4%
vs MPNet+21.4%+21.3%+25.1%

πŸš€ Quick Start

Installation

pip install transformers torch

Basic Usage

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:

# 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

ComponentPurposeHow It Helps
Temporal EmbeddingsEncodes memory age (today, this week, last month, etc.)"Yesterday" queries match recent memories
Memory Type EmbeddingsDistinguishes episodic/semantic/procedural"What do I like?" matches preferences, not events
Importance HeadAuxiliary task predicting memory priorityHelps 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

AspectDetails
Base Modelanswerdotai/ModernBERT-base (110M params)
Training Data500K synthetic agent memory samples
Memory DistributionEpisodic (40%), Semantic (35%), Procedural (25%)
Loss FunctionMultiple Negatives Ranking Loss + Importance MSE
Hard Negatives7 per sample (5 types: temporal, type confusion, topic drift, etc.)
Batch Size16-32 per GPU
Hardware2Γ— NVIDIA RTX 6000 Ada (48GB each)
Training Time~12 hours
PrecisionFP16 Mixed Precision
Final Val Loss0.877

πŸ—οΈ Use Cases

1. AI Agents with Long-Term Memory

# 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

# 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

# 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 CaseBest Model
AI agents with memoryβœ… AgentRank
Time-sensitive retrievalβœ… AgentRank
Conversational AIβœ… AgentRank
General document searchOpenAI / Cohere
Code searchCodeBERT
Scientific papersSciBERT

πŸ“ Model Family

ModelParametersSpeedQualityBest For
agentrank-small33M⚑⚑⚑ FastGoodReal-time agents, edge
agentrank-base149M⚑⚑ MediumBestQuality-critical apps
agentrank-reranker (coming)149M⚑ SlowerSuperiorTwo-stage retrieval

πŸ“š Citation

@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


πŸ“„ License

Apache 2.0 - Free for commercial use!


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

Built with ❀️ for the AI agent community

Contributors

vrushket

3 commits