pelinker.selection

Model-selection evaluation: load embeddings, draw stratified samples, evaluate screeners + clustering.

Production :meth:~pelinker.model.Linker.fit uses the same clustering subsample contract when LinkerFitConfig.clustering_sample_rows / base_seed / clustering_sample_index match the selection run.

evaluate_selection_from_paths(transform_config, optimization_config=None, *, file_path=None, file_paths=None, dfr=None, selected_labels=None, all_metrics_dfs=None, sample_index=0, embedding_read_status=None, show_embedding_read_progress=False)

Load, draw one stratified sample, and evaluate (convenience for single-shot callers).

Source code in pelinker/selection.py

def evaluate_selection_from_paths(
    transform_config: TransformConfig,
    optimization_config: ClusteringOptimizationConfig | None = None,
    *,
    file_path: pathlib.Path | None = None,
    file_paths: Sequence[pathlib.Path] | None = None,
    dfr: pd.DataFrame | None = None,
    selected_labels: set[str] | None = None,
    all_metrics_dfs: list[pd.DataFrame] | None = None,
    sample_index: int = 0,
    embedding_read_status: Callable[[str], None] | None = None,
    show_embedding_read_progress: bool = False,
) -> ModelSelectionReport | None:
    """
    Load, draw one stratified sample, and evaluate (convenience for single-shot callers).
    """
    config = optimization_config or ClusteringOptimizationConfig()
    base = load_selection_frame(
        file_path=file_path,
        file_paths=file_paths,
        dfr=dfr,
        config=config,
        selected_labels=selected_labels,
        embedding_read_status=embedding_read_status,
        show_embedding_read_progress=show_embedding_read_progress,
    )
    if base is None:
        return None
    sample_frame = draw_selection_sample(base, config, sample_index=sample_index)
    return evaluate_selection_sample(
        sample_frame,
        transform_config,
        optimization_config=config,
        all_metrics_dfs=all_metrics_dfs,
    )

evaluate_selection_sample(frame, transform_config, optimization_config=None, *, selected_labels=None, all_metrics_dfs=None, aggregation_level='mention', entities_pre_filtered=True)

Evaluate one selection draw: screener CV, transform, clustering grid, and HDBSCAN fit.

n_clusters_emergent in the returned report is at the grid-optimal best_size, not an arbitrary min_cluster_size. Compare with Linker.fit at a fixed MCS via :func:~pelinker.reporting.n_clusters_at_min_cluster_size on metrics_df.

When all_metrics_dfs is provided, appends this sample's grid metrics_df. Run :func:~pelinker.analysis.pooled_min_cluster_size_from_metrics_dfs after all bootstraps for a pooled min_cluster_size.

entities_pre_filtered=True skips min-mention-per-entity trimming (default when the frame came from :func:load_selection_frame).

Source code in pelinker/selection.py

def evaluate_selection_sample(
    frame: pd.DataFrame,
    transform_config: TransformConfig,
    optimization_config: ClusteringOptimizationConfig | None = None,
    *,
    selected_labels: set[str] | None = None,
    all_metrics_dfs: list[pd.DataFrame] | None = None,
    aggregation_level: Literal["mention", "entity"] = "mention",
    entities_pre_filtered: bool = True,
) -> ModelSelectionReport | None:
    """
    Evaluate one selection draw: screener CV, transform, clustering grid, and HDBSCAN fit.

    ``n_clusters_emergent`` in the returned report is at the grid-optimal ``best_size``,
    not an arbitrary ``min_cluster_size``. Compare with ``Linker.fit`` at a fixed MCS via
    :func:`~pelinker.reporting.n_clusters_at_min_cluster_size` on ``metrics_df``.

    When ``all_metrics_dfs`` is provided, appends this sample's grid ``metrics_df``. Run
    :func:`~pelinker.analysis.pooled_min_cluster_size_from_metrics_dfs` after all bootstraps
    for a pooled ``min_cluster_size``.

    ``entities_pre_filtered=True`` skips min-mention-per-entity trimming (default when the
    frame came from :func:`load_selection_frame`).
    """
    config = optimization_config or ClusteringOptimizationConfig()
    dfr = frame

    if "embed" not in dfr.columns or "entity" not in dfr.columns:
        return None

    if selected_labels is not None:
        dfr = dfr.loc[dfr["entity"].isin(selected_labels)].copy()
        if len(dfr) == 0:
            return None

    screener_cfg = config.ambient_screener
    neg_label = screener_cfg.negative_label

    if aggregation_level == "mention" and not entities_pre_filtered:
        if config.drop_rare_entities:
            dfr = drop_entities_with_few_mentions(
                dfr,
                config.min_mentions_per_entity,
                negative_label=neg_label,
            )
            if len(dfr) == 0:
                return None

    neg_mask, dfr_manifold = split_by_negative_label(dfr, neg_label)
    if len(dfr_manifold) == 0:
        return None

    number_properties = int(
        dfr_manifold["entity"].nunique()
    )  # for report metadata only

    tx_result = fit_transformer_on_manifold(dfr_manifold, transform_config)
    artifacts = score_transform_artifacts(
        dfr_manifold,
        tx_result.transformer,
        include_umap=True,
    )

    projection_cfg = config.projection_screener
    (
        X_embed_full,
        y_full,
        entity_full,
        orig_idx_full,
        X_m_cv,
        full_quality_artifacts,
    ) = _prepare_screener_cv_arrays(dfr, tx_result, projection_cfg, neg_label)

    unified = evaluate_all_screeners_cv(
        X_embed=X_embed_full,
        X_manifold=X_m_cv,
        y=y_full,
        entity=entity_full,
        orig_idx=orig_idx_full,
        screener_cfg=screener_cfg,
        oov_cfg=projection_cfg,
    )
    all_screener_cv: AllScreenerCvResult | None = None
    screener_oos_dp: PerDatapointScores | None = None
    if unified is not None:
        all_screener_cv, screener_oos_dp = unified

    umap_clustering_df = artifacts.umap_clustering_df().assign(
        entity=dfr_manifold["entity"].values
    )

    sizes = list(
        np.arange(
            config.resolved_min_scale(),
            config.max_scale,
            config.clustering_grid_step,
        )
    )
    metrics_df = evaluate_cluster_size_grid(
        umap_clustering_df,
        [c for c in umap_clustering_df.columns if c != "entity"],
        sizes,
    )
    if len(metrics_df) == 0:
        return None

    if all_metrics_dfs is not None:
        all_metrics_dfs.append(metrics_df)

    best_size, best_score = _solve_cluster_size_from_grid(metrics_df, config)

    cl_result = fit_hdbscan_on_umap(
        artifacts.umap_clustering,
        dfr_manifold,
        best_size,
        prediction_data=False,
    )
    labels = cl_result.cluster_labels
    fit_metrics = cl_result.fit_metrics
    assignments = _build_selection_assignments(dfr_manifold, labels)

    y_neg = entity_negative_label_mask_01(dfr_manifold["entity"], neg_label)
    mention_quality = mention_quality_frame(
        dfr,
        neg_mask=neg_mask,
        cluster_kb=labels,
        pca_residuals=full_quality_artifacts.pca_residuals,
        pca_mahalanobis=full_quality_artifacts.pca_mahalanobis,
        pca_spectral_entropy=full_quality_artifacts.pca_spectral_entropy,
        negative_label=neg_label,
    )

    return ModelSelectionReport(
        hyperparameters=ClusteringHyperparameters(min_cluster_size=best_size),
        best_score=best_score,
        number_properties=number_properties,
        n_clusters_emergent=fit_metrics.n_clusters_emergent,
        metrics_df=metrics_df,
        assignments=assignments,
        pca_residuals=artifacts.pca_residuals,
        pca_mahalanobis=artifacts.pca_mahalanobis,
        pca_spectral_entropy=artifacts.pca_spectral_entropy,
        oov_label=y_neg,
        umap_clustering=artifacts.umap_clustering,
        cluster_viz=artifacts.cluster_viz,
        cluster_viz_method=transform_config.cluster_viz_method,
        pca_reduced=artifacts.pca_reduced,
        all_screener_cv=all_screener_cv,
        screener_oos_datapoints=screener_oos_dp,
        ari=fit_metrics.ari,
        mention_quality=mention_quality,
    )

load_selection_frame(*, file_path=None, file_paths=None, dfr=None, config, selected_labels=None, embedding_read_status=None, show_embedding_read_progress=False)

Load and filter mention-level embeddings for model selection (no subsampling).

Provide exactly one of file_path, file_paths, or dfr.

Source code in pelinker/selection.py

def load_selection_frame(
    *,
    file_path: pathlib.Path | None = None,
    file_paths: Sequence[pathlib.Path] | None = None,
    dfr: pd.DataFrame | None = None,
    config: ClusteringOptimizationConfig,
    selected_labels: set[str] | None = None,
    embedding_read_status: Callable[[str], None] | None = None,
    show_embedding_read_progress: bool = False,
) -> pd.DataFrame | None:
    """
    Load and filter mention-level embeddings for model selection (no subsampling).

    Provide exactly one of ``file_path``, ``file_paths``, or ``dfr``.
    """
    sources = [file_path is not None, file_paths is not None, dfr is not None]
    if sum(bool(x) for x in sources) != 1:
        raise ValueError(
            "Provide exactly one of file_path=, file_paths=, or dfr= to load_selection_frame"
        )

    frame: pd.DataFrame | None
    if dfr is not None:
        frame = dfr.copy()
    else:
        if file_paths is not None:
            paths = list(file_paths)
        else:
            assert file_path is not None
            paths = [file_path]
        if len(paths) == 0:
            return None
        frame = concat_mention_level_embedding_sources(
            paths,
            batch_size=config.batch_size,
            read_status=embedding_read_status,
            show_read_progress=show_embedding_read_progress,
        )
        if frame is None or len(frame) == 0:
            return None

    if selected_labels is not None:
        frame = frame.loc[frame["entity"].isin(selected_labels)].copy()
        if len(frame) == 0:
            return None

    neg_label = config.ambient_screener.negative_label
    if config.drop_rare_entities:
        frame = drop_entities_with_few_mentions(
            frame,
            config.min_mentions_per_entity,
            negative_label=neg_label,
        )
        if len(frame) == 0:
            return None

    if config.max_mentions_per_entity is not None:
        frame = cap_mentions_per_entity(
            frame,
            max_mentions=config.max_mentions_per_entity,
            negative_label=neg_label,
            max_mentions_negative=config.max_mentions_negative,
            random_state=config.mention_cap_seed,
        )
        if len(frame) == 0:
            return None
    return frame