src/stage1.py

import numpy as np

def get_sparse_bins(sparse_results, cuts=(10, 25, 50, 100, 200, 1000)):
    bins = []
    prev = 0
    n = len(sparse_results)
    for cut in cuts:
        end = min(cut, n)
        bins.append({sparse_results[i][0] for i in range(prev, end)})
        prev = end
    while len(bins) < len(cuts):
        bins.append(set())
    return bins

def stage1_select(query_emb, sparse_results, cluster_manager, n_select=32):
    bins = get_sparse_bins(sparse_results)
    
    raw_scores = [s for _, s in sparse_results]
    min_s = min(raw_scores) if raw_scores else 0
    max_s = max(raw_scores) if raw_scores else 0
    norm_lut = {idx: ((score - min_s) / (max_s - min_s) if max_s > min_s else 0.0) for idx, score in sparse_results}

    cluster_scores = []
    v_bins = len(bins)
    
    for c_id, doc_idxs in cluster_manager.cluster_to_docs.items():
        doc_set = set(doc_idxs)
        P = [len(doc_set & bins[b]) for b in range(v_bins)]
        Q = [(float(np.mean([norm_lut[d] for d in (doc_set & bins[b])])) if (doc_set & bins[b]) else 0.0) for b in range(v_bins)]
        c_sim = float(np.dot(query_emb, cluster_manager.centroids[c_id]))
        cluster_scores.append((c_id, P, Q, c_sim))

    cluster_scores.sort(key=lambda x: (*x[1], x[3]), reverse=True)
    return cluster_scores[:n_select]