feat: add two-stage embedding retrieval and candidate ranking - #19
Conversation
📝 WalkthroughWalkthroughThe recommendation API now uses exact embedding retrieval to build a bounded candidate pool. ChangesRecommendation pipeline
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant RecommendationAPI
participant ExactEmbeddingRetriever
participant DeepSequenceModel
RecommendationAPI->>ExactEmbeddingRetriever: retrieve candidates with history, top_k, and exclusions
ExactEmbeddingRetriever-->>RecommendationAPI: return candidate IDs
RecommendationAPI->>DeepSequenceModel: rank candidate IDs
DeepSequenceModel-->>RecommendationAPI: return ranked recommendation IDs
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 71ce8cd91c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| query = model.embedding(history_ids[known_mask]).mean(dim=0) | ||
| query = F.normalize(query, dim=0, eps=1e-12) | ||
| catalogue = F.normalize( | ||
| model.embedding.weight[1 : model.num_items + 1], |
There was a problem hiding this comment.
When the catalogue exceeds the default 100-candidate pool, this cosine search can discard the model's highest-scoring items because training only optimizes the independent output_proj logits; it does not align averaged input embeddings with next-item retrieval. Moreover, src/training/train.py::model_predictions still evaluates and promotes bundles through full-catalogue model.recommend, so the quality metrics do not cover the path now used by the API. Evaluate the complete retrieval/ranking path for promotion or train an aligned retrieval objective before enabling this truncation.
Useful? React with 👍 / 👎.
| if not 1 <= top_k <= len(unique_candidates): | ||
| raise ValueError("top_k must not exceed the candidate pool") | ||
|
|
||
| logits = self.forward(item_seq) |
There was a problem hiding this comment.
For large catalogues, self.forward(item_seq) still computes output_proj logits for every item before the method gathers candidate scores, so the configured candidate-pool size does not bound ranking work. Since retrieval also scans the entire embedding table, this serving path adds another catalogue-linear pass without eliminating the original one and can increase latency or trigger the existing latency fallback. Compute the sequence context separately and apply only the selected rows of the projection to make candidate ranking actually bounded.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@app/core/model.py`:
- Around line 106-113: The rank_candidates flow in app/core/model.py (lines
106-113) should reuse the encoder context and apply the output projection only
to unique_candidates instead of calling forward to compute all item logits;
preserve candidate_scores and top_k behavior. Update
docs/architecture-walkthrough.md (lines 9-13) to describe the candidate-only
projection implementation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e52631b9-9fee-4031-a10f-3682d1970a05
📒 Files selected for processing (6)
app/api/routes.pyapp/core/config.pyapp/core/model.pyapp/core/retrieval.pydocs/architecture-walkthrough.mdtests/test_retrieval.py
| logits = self.forward(item_seq) | ||
| candidates = torch.tensor( | ||
| unique_candidates, | ||
| dtype=torch.long, | ||
| device=logits.device, | ||
| ) | ||
| candidate_scores = logits[0, candidates] | ||
| positions = torch.topk(candidate_scores, k=top_k).indices |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 5 'def forward|output_proj|def rank_candidates|candidate_scores' app/core/model.py
rg -n -C 4 'dense output projection|richer scoring capacity|candidate' \
docs/architecture-walkthrough.mdRepository: CoreyLeath-code/DeepSequence-Recommender
Length of output: 5293
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- model.py ---'
sed -n '1,155p' app/core/model.py
printf '%s\n' '--- rank_candidates call sites ---'
rg -n -C 5 'rank_candidates\(' app tests 2>/dev/null || true
printf '%s\n' '--- model tests and projection assumptions ---'
rg -n -C 4 'DeepSequenceModel|output_proj|rank_candidates|recommend\(' app tests 2>/dev/null || trueRepository: CoreyLeath-code/DeepSequence-Recommender
Length of output: 23908
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
model = ast.parse(Path("app/core/model.py").read_text(encoding="utf-8"))
routes = ast.parse(Path("app/api/routes.py").read_text(encoding="utf-8"))
def method(name):
for node in ast.walk(model):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name:
return node
raise SystemExit(f"missing method: {name}")
forward = method("forward")
rank = method("rank_candidates")
def calls(node, attr):
return [
n for n in ast.walk(node)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == attr
]
forward_projection = [
n for n in ast.walk(forward)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "output_proj"
]
rank_forward = calls(rank, "forward")
rank_candidate_index = [
n for n in ast.walk(rank)
if isinstance(n, ast.Subscript)
and isinstance(n.value, ast.Name)
and n.value.id == "logits"
]
route_rank = [
n for n in ast.walk(routes)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "rank_candidates"
]
print({
"forward_output_projection_calls": len(forward_projection),
"rank_candidates_forward_calls": len(rank_forward),
"rank_candidates_logits_subscripts": len(rank_candidate_index),
"api_rank_candidates_calls": len(route_rank),
"projection_output_width": "num_items + 1",
"candidate_validation": "1 <= candidate <= num_items",
})
PYRepository: CoreyLeath-code/DeepSequence-Recommender
Length of output: 424
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
model = ast.parse(Path("app/core/model.py").read_text(encoding="utf-8"))
routes = ast.parse(Path("app/api/routes.py").read_text(encoding="utf-8"))
classes = {
node.name: node
for node in model.body
if isinstance(node, ast.ClassDef)
}
deep_sequence = classes["DeepSequenceModel"]
methods = {
node.name: node
for node in deep_sequence.body
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
forward = methods["forward"]
rank = methods["rank_candidates"]
def attr_calls(node, name):
return [
n for n in ast.walk(node)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == name
]
projection_calls = [
n for n in ast.walk(forward)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "output_proj"
]
rank_logits_subscripts = [
n for n in ast.walk(rank)
if isinstance(n, ast.Subscript)
and isinstance(n.value, ast.Name)
and n.value.id == "logits"
]
route_rank_calls = [
n for n in ast.walk(routes)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "rank_candidates"
]
print({
"deep_sequence_forward_output_projection_calls": len(projection_calls),
"rank_candidates_forward_calls": len(attr_calls(rank, "forward")),
"rank_candidates_logits_subscripts": len(rank_logits_subscripts),
"api_rank_candidates_calls": len(route_rank_calls),
"projection_output_width": "num_items + 1",
"candidate_validation": "1 <= candidate <= num_items",
})
PYRepository: CoreyLeath-code/DeepSequence-Recommender
Length of output: 438
Restrict output projection to candidate IDs.
rank_candidates calls forward, which computes num_items + 1 logits before gathering candidates. Reuse the encoder context and project only the candidate rows. Update docs/architecture-walkthrough.md to match the implementation until this change is complete.
📍 Affects 2 files
app/core/model.py#L106-L113(this comment)docs/architecture-walkthrough.md#L9-L13
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@app/core/model.py` around lines 106 - 113, The rank_candidates flow in
app/core/model.py (lines 106-113) should reuse the encoder context and apply the
output projection only to unique_candidates instead of calling forward to
compute all item logits; preserve candidate_scores and top_k behavior. Update
docs/architecture-walkthrough.md (lines 9-13) to describe the candidate-only
projection implementation.
1 participant
Summary
Adds a real, backward-compatible two-stage serving boundary: exact embedding retrieval produces a bounded candidate pool, then the existing sequence model ranks only those candidates. It also adds an 8-minute diagram-based architecture walkthrough grounded in the implementation.
Architecture impact
ExactEmbeddingRetrieveruses normalized item-embedding similarity over the current in-memory catalogue and excludes already seen items.DeepSequenceModel.rank_candidatesruns the existing bidirectional-LSTM attention model and gathers scores only for the retrieval result.Files changed
app/core/retrieval.py— new exact, bounded candidate-retrieval contract.app/core/model.py— candidate-only ranking method, preserving the existing full-cataloguerecommendmethod for training compatibility.app/api/routes.py— calls retrieval before candidate ranking.app/core/config.py— addsRETRIEVAL_CANDIDATE_POOL_SIZEwith a safe default.tests/test_retrieval.py— validates seen-item exclusion, bounded candidate ordering, candidate-only ranking, and impossible requests.docs/architecture-walkthrough.md— 8-minute diagram/script explaining two-stage retrieval vs. ranking, the absent reasoning boundary, and tradeoffs.Validation performed
main; the diff is limited to the retrieval/ranking boundary, config, focused tests, and implementation-grounded walkthrough.ghcredentials are stale.Risks
RETRIEVAL_CANDIDATE_POOL_SIZEmust be calibrated against a leakage-safe evaluation set before a production configuration change.Follow-up recommendations
Refs #13, #16
Summary by CodeRabbit