diff options
| author | YurenHao0426 <blackhao0426@gmail.com> | 2026-01-27 09:57:37 -0600 |
|---|---|---|
| committer | YurenHao0426 <blackhao0426@gmail.com> | 2026-01-27 09:57:37 -0600 |
| commit | dc801c07cf38b0c495686463e6ca6f871a64440e (patch) | |
| tree | 599f03114775921dbc472403c701f4a3a8ea188a /src/personalization/retrieval/pipeline.py | |
| parent | e43b3f8aa36c198b95c1e46bea2eaf3893b13dc3 (diff) | |
Add collaborativeagents module and update gitignore
- Add collaborativeagents subproject with adapters, agents, and evaluation modules
- Update .gitignore to exclude large binary files (.whl, .tar), wandb logs, and results
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Diffstat (limited to 'src/personalization/retrieval/pipeline.py')
| -rw-r--r-- | src/personalization/retrieval/pipeline.py | 42 |
1 files changed, 26 insertions, 16 deletions
diff --git a/src/personalization/retrieval/pipeline.py b/src/personalization/retrieval/pipeline.py index 3d3eeb7..e83940d 100644 --- a/src/personalization/retrieval/pipeline.py +++ b/src/personalization/retrieval/pipeline.py @@ -110,10 +110,14 @@ def retrieve_with_policy( candidates = [memory_cards[i] for i in dense_idx] candidate_docs = [c.note_text for c in candidates] - + # 2. Rerank base score (P(yes|q,m)) - base_scores = np.array(reranker.score(query, candidate_docs)) - + # Skip reranking if we have fewer candidates than topk_rerank (saves GPU memory) + if len(candidates) <= topk_rerank: + base_scores = np.ones(len(candidates)) # Uniform scores + else: + base_scores = np.array(reranker.score(query, candidate_docs)) + # 3. Policy Scoring (Softmax) user_state: UserState = user_store.get_state(user_id) candidate_vectors = item_vectors[dense_idx] # [K, k] @@ -181,29 +185,35 @@ def retrieve_no_policy( # 1. Dense retrieval dense_idx = dense_topk_indices( - query, - embed_model, - memory_embeddings, + query, + embed_model, + memory_embeddings, valid_indices=valid_indices, topk=topk_dense ) - + if not dense_idx: return [], np.array([]), np.array([]), [], np.array([]) candidates = [memory_cards[i] for i in dense_idx] candidate_docs = [c.note_text for c in candidates] - + # 2. Rerank base score (P(yes|q,m)) - base_scores = np.array(reranker.score(query, candidate_docs)) - - # 3. Deterministic Top-K selection based on rerank scores ONLY (no policy) - k = min(topk_rerank, len(base_scores)) - top_indices_local = base_scores.argsort()[-k:][::-1] - chosen_indices = top_indices_local.tolist() - + # Skip reranking if we have fewer candidates than topk_rerank (saves GPU memory) + if len(candidates) <= topk_rerank: + # Just return all candidates without reranking + base_scores = np.ones(len(candidates)) # Uniform scores + chosen_indices = list(range(len(candidates))) + else: + base_scores = np.array(reranker.score(query, candidate_docs)) + + # 3. Deterministic Top-K selection based on rerank scores ONLY (no policy) + k = min(topk_rerank, len(base_scores)) + top_indices_local = base_scores.argsort()[-k:][::-1] + chosen_indices = top_indices_local.tolist() + # Get scores for chosen items (for logging compatibility) - chosen_scores = base_scores[top_indices_local] + chosen_scores = base_scores[chosen_indices] # Return empty item vectors (not used in NoPersonal mode) # Return rerank scores as the "probs" field for logging compatibility |
