Retrieval

Retrieval#

Stage 1 of the recommendation pipeline. Retrievers generate candidate items from large catalogs efficiently, prioritizing recall over precision.

Quick Start#

from corerec.retrieval import SemanticRetriever, CollaborativeRetriever, EnsembleRetriever

retriever = SemanticRetriever(encoder="all-MiniLM-L6-v2")
retriever.fit(items=item_catalog)
candidates = retriever.retrieve(query=user_id, top_k=100)

API Reference#

Retrieval Module - Candidate Generation (Stage 1)

This module provides retrievers that generate candidate items from large catalogs efficiently. Retrievers prioritize recall over precision.

Usage:

from corerec.retrieval import SemanticRetriever, CollaborativeRetriever

retriever = SemanticRetriever(encoder=”all-MiniLM-L6-v2”) retriever.index(items) candidates = retriever.retrieve(query, top_k=100)

class corerec.retrieval.BaseRetriever(name: str | None = None)[source]

Bases: ABC

Abstract base class for all retrievers.

A retriever takes a query (user, context, or text) and returns a set of candidate items. Subclasses implement different retrieval strategies: collaborative filtering, semantic search, popularity-based, etc.

Design goals: - Fast: O(log n) or O(1) lookup, not O(n) - Recall-focused: better to include irrelevant than miss relevant - Composable: multiple retrievers can be ensembled

abstractmethod fit(**kwargs) → BaseRetriever[source]

Fit the retriever on training data.

What “fit” means depends on the retriever type: - CollaborativeRetriever: build item-item similarity matrix - SemanticRetriever: encode items into embeddings - PopularityRetriever: count item frequencies

Returns self for method chaining.

property is_fitted: bool

abstractmethod retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]

Retrieve candidate items for a query.

Parameters:

query – the query to retrieve for. Can be: - user_id for collaborative retrieval - text string for semantic retrieval - embedding vector for ANN search
top_k – number of candidates to return
**kwargs – retriever-specific parameters

Returns:

RetrievalResult containing ranked candidates

retrieve_batch(queries: List[Any], top_k: int = 100, **kwargs) → List[RetrievalResult][source]: Retrieve for multiple queries. Default impl just loops. Subclasses can override for batched efficiency.

class corerec.retrieval.Candidate(item_id: ~typing.Any, score: float, source: str = 'unknown', metadata: ~typing.Dict[str, ~typing.Any] = <factory>)[source]

Bases: object

A single candidate item from retrieval.

source: str = 'unknown': which retriever produced this candidate

item_id: Any: unique identifier for the item

score: float: retrieval score (higher = more relevant)

metadata: Dict[str, Any]: optional extra info (embeddings, features, etc)

class corerec.retrieval.RetrievalResult(candidates: List[Candidate], query_id: Any = None, retriever_name: str = 'unknown', timing_ms: float = 0.0)[source]

Bases: object

Result from a retrieval operation.

item_ids() → List[Any][source]: Extract just the item IDs.

query_id: Any = None: identifier for the query (user_id, session_id, etc)

retriever_name: str = 'unknown': name of the retriever that produced this

scores() → ndarray[source]: Extract scores as numpy array.

timing_ms: float = 0.0: how long retrieval took in milliseconds

top_k(k: int) → List[Candidate][source]: Return top k candidates by score.

candidates: List[Candidate]: list of retrieved candidates

class corerec.retrieval.CollaborativeRetriever(model: Any | None = None, name: str = 'collaborative', exclude_seen: bool = True)[source]

Bases: BaseRetriever

Retriever based on collaborative filtering signals.

This wraps any CF model that has a recommend() method, making it usable in retrieval pipelines. The CF model handles the actual similarity computation; this class just adapts the interface.

Example

from corerec.collaborative import SAR from corerec.retrieval import CollaborativeRetriever

sar = SAR(similarity_type=’jaccard’) sar.fit(train_df)

retriever = CollaborativeRetriever(model=sar) candidates = retriever.retrieve(user_id=123, top_k=100)

__init__(model: Any | None = None, name: str = 'collaborative', exclude_seen: bool = True)[source]

Parameters:

model – a fitted CF model with recommend() method
name – identifier for this retriever
exclude_seen – whether to exclude items user has seen

fit(model: Any | None = None, **fit_kwargs) → CollaborativeRetriever[source]

Fit the underlying CF model, or attach a pre-fitted model.

Parameters:

model – CF model to use (if not set in __init__)
**fit_kwargs – passed to model.fit() if model needs fitting

retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]

Retrieve candidates for a user.

Parameters:

query – user_id to retrieve for
top_k – number of candidates

Returns:

RetrievalResult with candidates

retrieve_batch(queries: List[Any], top_k: int = 100, **kwargs) → List[RetrievalResult][source]: Batch retrieval - uses model’s batch method if available.

class corerec.retrieval.SemanticRetriever(encoder: str | Any | None = None, embedding_dim: int = 384, similarity: str = 'cosine', name: str = 'semantic')[source]

Bases: BaseRetriever

Retriever based on semantic similarity of text embeddings.

Encodes item descriptions (or any text) into dense vectors, builds an index, and retrieves via approximate nearest neighbor search.

Example:

retriever = SemanticRetriever(
    encoder="sentence-transformers/all-MiniLM-L6-v2"
)
retriever.fit(
    item_ids=[1, 2, 3],
    item_texts=["red shoes", "blue jacket", "green hat"]
)

# text query
results = retriever.retrieve("comfortable footwear", top_k=10)

# or embed a user profile as query
user_embedding = retriever.encode("likes outdoor hiking gear")
results = retriever.retrieve(user_embedding, top_k=10)

__init__(encoder: str | Any | None = None, embedding_dim: int = 384, similarity: str = 'cosine', name: str = 'semantic')[source]

Parameters:

encoder – sentence transformer model name or encoder instance. If string, will try to load from sentence-transformers. If None, you must provide embeddings directly.
embedding_dim – dimension of embeddings (auto-detected if encoder provided)
similarity – ‘cosine’ or ‘dot’ for similarity computation
name – identifier for this retriever

encode(texts: str | List[str]) → ndarray[source]

Encode text(s) into embeddings.

Parameters:: texts – single string or list of strings
Returns:: embeddings array of shape (n, dim) or (dim,) for single text

property encoder: Lazy load encoder on first access.

fit(item_ids: List[Any], item_texts: List[str] | None = None, item_embeddings: ndarray | None = None, **kwargs) → SemanticRetriever[source]

Index items for retrieval.

Parameters:

item_ids – unique identifiers for each item
item_texts – text descriptions to encode (if not providing embeddings)
item_embeddings – pre-computed embeddings (if not providing texts)

You must provide either item_texts or item_embeddings.

retrieve(query: str | ndarray, top_k: int = 100, **kwargs) → RetrievalResult[source]

Retrieve items similar to query.

Parameters:

query – text string or embedding vector
top_k – number of candidates to return

Returns:

RetrievalResult with candidates sorted by similarity

retrieve_batch(queries: List[str | ndarray], top_k: int = 100, **kwargs) → List[RetrievalResult][source]: Batch retrieval - more efficient for multiple queries.

class corerec.retrieval.PopularityRetriever(name: str = 'popularity', time_decay: float | None = None)[source]

Bases: BaseRetriever

Retriever that returns most popular items.

“Popular” can mean different things: - Most interactions (views, clicks, purchases) - Highest average rating - Most recent trending

This is a simple but effective baseline, especially for: - Cold-start users with no history - Fallback when other retrievers fail - Diversity injection in ensembles

Example

retriever = PopularityRetriever() retriever.fit(item_ids, interaction_counts) candidates = retriever.retrieve(user_id=None, top_k=50)

__init__(name: str = 'popularity', time_decay: float | None = None)[source]

Parameters:

name – identifier for this retriever
time_decay – if set, apply exponential decay based on recency

fit(item_ids: List[Any], scores: List[float] | None = None, interaction_counts: List[int] | None = None, timestamps: List[float] | None = None, **kwargs) → PopularityRetriever[source]

Compute popularity scores for items.

Parameters:

item_ids – unique identifiers for items
scores – pre-computed popularity scores (if available)
interaction_counts – raw interaction counts to use as popularity
timestamps – if provided with time_decay, applies recency weighting

Provide either scores or interaction_counts.

get_item_popularity(item_id: Any) → float[source]: Get popularity score for a specific item.

retrieve(query: Any = None, top_k: int = 100, exclude_items: List[Any] | None = None, **kwargs) → RetrievalResult[source]

Retrieve most popular items.

Parameters:

query – ignored (popularity is query-independent)
top_k – number of items to return
exclude_items – items to exclude from results

Returns:

RetrievalResult with most popular items

class corerec.retrieval.VectorIndex(index_type: str = 'flat', metric: str = 'cosine', nlist: int = 100, nprobe: int = 10, name: str = 'vector_index')[source]

Bases: BaseRetriever

High-performance vector similarity search using ANN algorithms.

Wraps FAISS (or falls back to brute-force numpy) for scalable nearest neighbor search on embeddings.

For catalogs with >100K items, this is essential for sub-millisecond retrieval. Smaller catalogs can use brute-force without issues.

Example

index = VectorIndex(index_type=”flat”, metric=”cosine”) index.fit(item_ids, embeddings)

# single query results = index.retrieve(query_embedding, top_k=100)

# batch results = index.retrieve_batch(query_embeddings, top_k=100)

__init__(index_type: str = 'flat', metric: str = 'cosine', nlist: int = 100, nprobe: int = 10, name: str = 'vector_index')[source]

Parameters:

index_type – ‘flat’ (exact), ‘ivf’ (approximate), or ‘hnsw’
metric – ‘cosine’, ‘l2’, or ‘ip’ (inner product)
nlist – number of clusters for IVF index
nprobe – number of clusters to search for IVF
name – identifier for this retriever

fit(item_ids: List[Any], embeddings: ndarray, **kwargs) → VectorIndex[source]

Build the vector index.

Parameters:

item_ids – unique identifiers for each embedding
embeddings – array of shape (n_items, dim)

retrieve(query: ndarray, top_k: int = 100, **kwargs) → RetrievalResult[source]

Search for nearest neighbors.

Parameters:

query – query embedding vector
top_k – number of neighbors to return

retrieve_batch(queries: List[ndarray], top_k: int = 100, **kwargs) → List[RetrievalResult][source]: Batch search - more efficient than individual queries.

class corerec.retrieval.EnsembleRetriever(retrievers: List[Tuple[str, BaseRetriever, float]], strategy: str = 'rrf', k_per_retriever: int | None = None, name: str = 'ensemble')[source]

Bases: BaseRetriever

Combines candidates from multiple retrievers.

Different retrievers capture different signals: - Collaborative: “users like you bought X” - Semantic: “X matches your query meaning” - Popularity: “X is trending”

Ensembling these improves recall and handles cases where one retriever might fail (cold start, sparse history, etc).

Fusion strategies: - union – combine all candidates, dedupe, keep max score - rrf – Reciprocal Rank Fusion (position-based, score-agnostic) - weighted – weighted sum of normalized scores

Example:

ensemble = EnsembleRetriever(
    retrievers=[
        ("collab", collaborative_retriever, 1.0),
        ("semantic", semantic_retriever, 0.5),
        ("popular", popularity_retriever, 0.3),
    ],
    strategy="rrf"
)
ensemble.fit()  # fits underlying retrievers if needed
candidates = ensemble.retrieve(query, top_k=100)

__init__(retrievers: List[Tuple[str, BaseRetriever, float]], strategy: str = 'rrf', k_per_retriever: int | None = None, name: str = 'ensemble')[source]

Parameters:

retrievers – list of (name, retriever, weight) tuples
strategy – ‘union’, ‘rrf’, or ‘weighted’
k_per_retriever – candidates to get from each retriever (defaults to 2x final top_k)
name – identifier for this retriever

fit(**kwargs) → EnsembleRetriever[source]

Fit all underlying retrievers.

Pass retriever-specific kwargs as retriever_name__param. e.g., collab__model=sar_model

retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]: Retrieve and merge candidates from all sources.

Base Classes#

Base classes for retrieval stage.

Retrievers are responsible for quickly narrowing down millions of items to hundreds of candidates. Speed matters more than precision here.

class corerec.retrieval.base.Candidate(item_id: ~typing.Any, score: float, source: str = 'unknown', metadata: ~typing.Dict[str, ~typing.Any] = <factory>)[source]#

Bases: object

A single candidate item from retrieval.

item_id: Any#: unique identifier for the item

score: float#: retrieval score (higher = more relevant)

source: str = 'unknown'#: which retriever produced this candidate

metadata: Dict[str, Any]#: optional extra info (embeddings, features, etc)

class corerec.retrieval.base.RetrievalResult(candidates: List[Candidate], query_id: Any = None, retriever_name: str = 'unknown', timing_ms: float = 0.0)[source]#

Bases: object

Result from a retrieval operation.

candidates: List[Candidate]#: list of retrieved candidates

query_id: Any = None#: identifier for the query (user_id, session_id, etc)

retriever_name: str = 'unknown'#: name of the retriever that produced this

timing_ms: float = 0.0#: how long retrieval took in milliseconds

top_k(k: int) → List[Candidate][source]#: Return top k candidates by score.

item_ids() → List[Any][source]#: Extract just the item IDs.

scores() → ndarray[source]#: Extract scores as numpy array.

class corerec.retrieval.base.BaseRetriever(name: str | None = None)[source]#

Bases: ABC

Abstract base class for all retrievers.

A retriever takes a query (user, context, or text) and returns a set of candidate items. Subclasses implement different retrieval strategies: collaborative filtering, semantic search, popularity-based, etc.

Design goals: - Fast: O(log n) or O(1) lookup, not O(n) - Recall-focused: better to include irrelevant than miss relevant - Composable: multiple retrievers can be ensembled

property is_fitted: bool#

abstractmethod fit(**kwargs) → BaseRetriever[source]#

Fit the retriever on training data.

What “fit” means depends on the retriever type: - CollaborativeRetriever: build item-item similarity matrix - SemanticRetriever: encode items into embeddings - PopularityRetriever: count item frequencies

Returns self for method chaining.

abstractmethod retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]#

Retrieve candidate items for a query.

Parameters:

query – the query to retrieve for. Can be: - user_id for collaborative retrieval - text string for semantic retrieval - embedding vector for ANN search
top_k – number of candidates to return
**kwargs – retriever-specific parameters

Returns:

RetrievalResult containing ranked candidates

retrieve_batch(queries: List[Any], top_k: int = 100, **kwargs) → List[RetrievalResult][source]#: Retrieve for multiple queries. Default impl just loops. Subclasses can override for batched efficiency.

Implementations#

Collaborative Retriever#

Collaborative Filtering based Retrieval

Wraps traditional CF algorithms (SAR, ALS, etc) as retrievers for use in multi-stage pipelines.

class corerec.retrieval.collaborative.CollaborativeRetriever(model: Any | None = None, name: str = 'collaborative', exclude_seen: bool = True)[source]#

Bases: BaseRetriever

Retriever based on collaborative filtering signals.

This wraps any CF model that has a recommend() method, making it usable in retrieval pipelines. The CF model handles the actual similarity computation; this class just adapts the interface.

Example

from corerec.collaborative import SAR from corerec.retrieval import CollaborativeRetriever

sar = SAR(similarity_type=’jaccard’) sar.fit(train_df)

retriever = CollaborativeRetriever(model=sar) candidates = retriever.retrieve(user_id=123, top_k=100)

__init__(model: Any | None = None, name: str = 'collaborative', exclude_seen: bool = True)[source]#

Parameters:

model – a fitted CF model with recommend() method
name – identifier for this retriever
exclude_seen – whether to exclude items user has seen

fit(model: Any | None = None, **fit_kwargs) → CollaborativeRetriever[source]#

Fit the underlying CF model, or attach a pre-fitted model.

Parameters:

model – CF model to use (if not set in __init__)
**fit_kwargs – passed to model.fit() if model needs fitting

retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]#

Retrieve candidates for a user.

Parameters:

query – user_id to retrieve for
top_k – number of candidates

Returns:

RetrievalResult with candidates

retrieve_batch(queries: List[Any], top_k: int = 100, **kwargs) → List[RetrievalResult][source]#: Batch retrieval - uses model’s batch method if available.

Semantic Retriever#

Semantic Retrieval using Text Embeddings

Uses sentence transformers or similar encoders to embed items and queries into a shared vector space, then retrieves via ANN search.

class corerec.retrieval.semantic.SemanticRetriever(encoder: str | Any | None = None, embedding_dim: int = 384, similarity: str = 'cosine', name: str = 'semantic')[source]#

Bases: BaseRetriever

Retriever based on semantic similarity of text embeddings.

Encodes item descriptions (or any text) into dense vectors, builds an index, and retrieves via approximate nearest neighbor search.

Example:

retriever = SemanticRetriever(
    encoder="sentence-transformers/all-MiniLM-L6-v2"
)
retriever.fit(
    item_ids=[1, 2, 3],
    item_texts=["red shoes", "blue jacket", "green hat"]
)

# text query
results = retriever.retrieve("comfortable footwear", top_k=10)

# or embed a user profile as query
user_embedding = retriever.encode("likes outdoor hiking gear")
results = retriever.retrieve(user_embedding, top_k=10)

__init__(encoder: str | Any | None = None, embedding_dim: int = 384, similarity: str = 'cosine', name: str = 'semantic')[source]#

Parameters:

encoder – sentence transformer model name or encoder instance. If string, will try to load from sentence-transformers. If None, you must provide embeddings directly.
embedding_dim – dimension of embeddings (auto-detected if encoder provided)
similarity – ‘cosine’ or ‘dot’ for similarity computation
name – identifier for this retriever

property encoder#: Lazy load encoder on first access.

encode(texts: str | List[str]) → ndarray[source]#

Encode text(s) into embeddings.

Parameters:: texts – single string or list of strings
Returns:: embeddings array of shape (n, dim) or (dim,) for single text

fit(item_ids: List[Any], item_texts: List[str] | None = None, item_embeddings: ndarray | None = None, **kwargs) → SemanticRetriever[source]#

Index items for retrieval.

Parameters:

item_ids – unique identifiers for each item
item_texts – text descriptions to encode (if not providing embeddings)
item_embeddings – pre-computed embeddings (if not providing texts)

You must provide either item_texts or item_embeddings.

retrieve(query: str | ndarray, top_k: int = 100, **kwargs) → RetrievalResult[source]#

Retrieve items similar to query.

Parameters:

query – text string or embedding vector
top_k – number of candidates to return

Returns:

RetrievalResult with candidates sorted by similarity

retrieve_batch(queries: List[str | ndarray], top_k: int = 100, **kwargs) → List[RetrievalResult][source]#: Batch retrieval - more efficient for multiple queries.

Popularity Retriever#

Popularity-based Retrieval

Simple baseline that returns most popular items. Useful for cold-start users and as a component in ensembles.

class corerec.retrieval.popularity.PopularityRetriever(name: str = 'popularity', time_decay: float | None = None)[source]#

Bases: BaseRetriever

Retriever that returns most popular items.

“Popular” can mean different things: - Most interactions (views, clicks, purchases) - Highest average rating - Most recent trending

This is a simple but effective baseline, especially for: - Cold-start users with no history - Fallback when other retrievers fail - Diversity injection in ensembles

Example

retriever = PopularityRetriever() retriever.fit(item_ids, interaction_counts) candidates = retriever.retrieve(user_id=None, top_k=50)

__init__(name: str = 'popularity', time_decay: float | None = None)[source]#

Parameters:

name – identifier for this retriever
time_decay – if set, apply exponential decay based on recency

fit(item_ids: List[Any], scores: List[float] | None = None, interaction_counts: List[int] | None = None, timestamps: List[float] | None = None, **kwargs) → PopularityRetriever[source]#

Compute popularity scores for items.

Parameters:

item_ids – unique identifiers for items
scores – pre-computed popularity scores (if available)
interaction_counts – raw interaction counts to use as popularity
timestamps – if provided with time_decay, applies recency weighting

Provide either scores or interaction_counts.

retrieve(query: Any = None, top_k: int = 100, exclude_items: List[Any] | None = None, **kwargs) → RetrievalResult[source]#

Retrieve most popular items.

Parameters:

query – ignored (popularity is query-independent)
top_k – number of items to return
exclude_items – items to exclude from results

Returns:

RetrievalResult with most popular items

get_item_popularity(item_id: Any) → float[source]#: Get popularity score for a specific item.

Ensemble Retriever#

Ensemble Retriever - Combines Multiple Retrieval Sources

In production systems, using multiple retrieval channels improves both coverage and quality. This module provides strategies for merging results from different retrievers.

class corerec.retrieval.ensemble.EnsembleRetriever(retrievers: List[Tuple[str, BaseRetriever, float]], strategy: str = 'rrf', k_per_retriever: int | None = None, name: str = 'ensemble')[source]#

Bases: BaseRetriever

Combines candidates from multiple retrievers.

Different retrievers capture different signals: - Collaborative: “users like you bought X” - Semantic: “X matches your query meaning” - Popularity: “X is trending”

Ensembling these improves recall and handles cases where one retriever might fail (cold start, sparse history, etc).

Fusion strategies: - union – combine all candidates, dedupe, keep max score - rrf – Reciprocal Rank Fusion (position-based, score-agnostic) - weighted – weighted sum of normalized scores

Example:

ensemble = EnsembleRetriever(
    retrievers=[
        ("collab", collaborative_retriever, 1.0),
        ("semantic", semantic_retriever, 0.5),
        ("popular", popularity_retriever, 0.3),
    ],
    strategy="rrf"
)
ensemble.fit()  # fits underlying retrievers if needed
candidates = ensemble.retrieve(query, top_k=100)

__init__(retrievers: List[Tuple[str, BaseRetriever, float]], strategy: str = 'rrf', k_per_retriever: int | None = None, name: str = 'ensemble')[source]#

Parameters:

retrievers – list of (name, retriever, weight) tuples
strategy – ‘union’, ‘rrf’, or ‘weighted’
k_per_retriever – candidates to get from each retriever (defaults to 2x final top_k)
name – identifier for this retriever

fit(**kwargs) → EnsembleRetriever[source]#

Fit all underlying retrievers.

Pass retriever-specific kwargs as retriever_name__param. e.g., collab__model=sar_model

retrieve(query: Any, top_k: int = 100, **kwargs) → RetrievalResult[source]#: Retrieve and merge candidates from all sources.

Vector Index#

Vector Index for Approximate Nearest Neighbor Search

Wraps FAISS or other ANN libraries for efficient similarity search on large embedding collections.

class corerec.retrieval.index.VectorIndex(index_type: str = 'flat', metric: str = 'cosine', nlist: int = 100, nprobe: int = 10, name: str = 'vector_index')[source]#

Bases: BaseRetriever

High-performance vector similarity search using ANN algorithms.

Wraps FAISS (or falls back to brute-force numpy) for scalable nearest neighbor search on embeddings.

For catalogs with >100K items, this is essential for sub-millisecond retrieval. Smaller catalogs can use brute-force without issues.

Example

index = VectorIndex(index_type=”flat”, metric=”cosine”) index.fit(item_ids, embeddings)

# single query results = index.retrieve(query_embedding, top_k=100)

# batch results = index.retrieve_batch(query_embeddings, top_k=100)

__init__(index_type: str = 'flat', metric: str = 'cosine', nlist: int = 100, nprobe: int = 10, name: str = 'vector_index')[source]#

Parameters:

index_type – ‘flat’ (exact), ‘ivf’ (approximate), or ‘hnsw’
metric – ‘cosine’, ‘l2’, or ‘ip’ (inner product)
nlist – number of clusters for IVF index
nprobe – number of clusters to search for IVF
name – identifier for this retriever

fit(item_ids: List[Any], embeddings: ndarray, **kwargs) → VectorIndex[source]#

Build the vector index.

Parameters:

item_ids – unique identifiers for each embedding
embeddings – array of shape (n_items, dim)

retrieve(query: ndarray, top_k: int = 100, **kwargs) → RetrievalResult[source]#

Search for nearest neighbors.

Parameters:

query – query embedding vector
top_k – number of neighbors to return

retrieve_batch(queries: List[ndarray], top_k: int = 100, **kwargs) → List[RetrievalResult][source]#: Batch search - more efficient than individual queries.

Retrieval

Contents

Retrieval#

Quick Start#

API Reference#

Base Classes#

Implementations#

Collaborative Retriever#

Semantic Retriever#

Popularity Retriever#

Ensemble Retriever#

Vector Index#