def cosine_sim_matrix(query_vec, doc_matrix):
    q = query_vec / (np.linalg.norm(query_vec) + 1e-12)
    d = doc_matrix / np.linalg.norm(doc_matrix, axis=1, keepdims=True)
    return d @ q

def retrieve_pages(page_df, line_df, question, top_k=3):
    q_vec = np.asarray(get_embedding(question), dtype=np.float32)
    doc_matrix = np.vstack(page_df["embedding"].values)
    sims = cosine_sim_matrix(q_vec, doc_matrix)

    scored = page_df.copy()
    scored["similarity"] = sims
    retrieved_pages_df = scored.nlargest(top_k, "similarity")

    kept_pages = retrieved_pages_df["page_num"].tolist()
    filtered_line_df = line_df[line_df["page_num"].isin(kept_pages)]
    return retrieved_pages_df, filtered_line_df