Function rerank_episodes

experiments/baselines/knowledge_entity/runner.py:733–814

(
    rows: List[Dict[str, Any]],
    label_rows: List[Dict[str, Any]],
    *,
    top_k: int,
    roles_file: Path,
    input_dir: Path,
)

def rerank_episodes(
    rows: List[Dict[str, Any]],
    label_rows: List[Dict[str, Any]],
    *,
    top_k: int,
    roles_file: Path,
    input_dir: Path,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    roles = load_roles(roles_file)
    user_metadata = load_user_metadata(input_dir)
    states = build_user_states(user_metadata, roles, rows)
    label_map = build_label_map(label_rows)

    output_rows: List[Dict[str, Any]] = []
    grouped = group_by_episode(rows)
    method_stats = {
        "method_key": METHOD_KEY,
        "method_name": METHOD_NAME,
        "episodes": 0,
        "top_k": top_k,
        "uses_dynamic_profile": False,
        "uses_reading_reports": False,
        "uses_live_entity_api": False,
        "users": {},
    }

    for episode_id, episode_rows in sorted(grouped.items(), key=episode_sort_key):
        if not episode_rows:
            continue
        user_id = str(episode_rows[0].get("user_id") or "").strip()
        if user_id not in states:
            states[user_id] = UserKnowledgeEntityState(profile_text="", profile_entities=[])
        state = states[user_id]

        entity_texts = [entities_text(extract_entities(row)) for row in episode_rows]
        day_texts = [paper_text(row) for row in episode_rows]
        idf = build_idf_from_texts([
            state.profile_text,
            entities_text(state.profile_entities),
            *state.feedback_texts(),
            *entity_texts,
            *day_texts,
        ])
        state.prepare_vectors(idf)
        max_log_citations = max_log_citations_for_day(episode_rows)

        scored_rows: List[Tuple[float, str, Dict[str, Any], Dict[str, Any]]] = []
        for row in episode_rows:
            identity = paper_identity(row)
            rep = build_paper_representation(row, idf)
            score_payload = score_row(
                row,
                rep,
                state,
                max_log_citations=max_log_citations,
            )
            scored_rows.append((score_payload["system_score"], identity, row, score_payload))