pelinker.cli.fit

FitCliConfig dataclass

Hydra config for python -m pelinker.cli.fit.

Source code in pelinker/cli/fit.py

@dataclass
class FitCliConfig:
    """Hydra config for ``python -m pelinker.cli.fit``."""

    model_type: str = "pubmedbert"
    layers_spec: str = "1"
    kb_path: str = MISSING
    selection_report: str | None = None
    """``selected_hyperparameters.json`` (or the report dir containing it) from
    ``pelinker-model-selection`` / ``pelinker-dim-selection``.

    Fills in ``pca_components``, ``umap_dim``, ``umap_n_neighbors`` and
    ``min_cluster_size`` when those are not set explicitly here. Explicit overrides always
    win, and any disagreement is logged rather than silently resolved."""
    pca_components: int | None = None
    """PCA components; omit to take the selection report's value, else 100."""
    umap_dim: int | None = None
    """UMAP output dimension; omit to take the selection report's value, else 8."""
    umap_n_neighbors: int | None = None
    """UMAP ``n_neighbors``; omit for the library default (15). Scale-dependent — see
    ``pelinker-scale-curve`` when the fit N differs markedly from the selection N."""
    cluster_viz_method: str = "pca"
    drop_rare_entities: bool = False
    min_mentions_per_entity: int = 20
    max_mentions_per_entity: int | None = None
    max_mentions_negative: int | None = None
    mention_cap_seed: int | None = None
    """Seed for per-entity mention cap; defaults to ``seed`` when omitted."""
    seed: int = 13
    """Bootstrap seed for clustering subsample draws (``base_seed``); also default for mention-cap and screener draws."""
    pca_seed: int = 13
    """Random seed for PCA and cluster-viz PCA."""
    umap_seed: int | None = None
    """UMAP random seed; omit (None) for parallel UMAP. Set for reproducible production fits."""
    clustering_sample_rows: int | None = None
    """Max mention rows per clustering bootstrap draw (stratified). None = use all loaded rows."""
    clustering_sample_index: int = 0
    """Bootstrap index for clustering subsample (match model-selection ``sample_idx``)."""
    # Stage-B HDBSCAN ``min_cluster_size`` (choose upstream, e.g. ``pelinker.model_selection``).
    min_cluster_size: int | None = None
    """Explicit HDBSCAN ``min_cluster_size``. Omit to resolve from ``scale_curve_path``,
    or fall back to 20 when neither is given. An explicit value always wins."""
    scale_curve_path: str | None = None
    """``scale_curve.json`` from ``pelinker-scale-curve``. When set (and
    ``min_cluster_size`` is not), ``min_cluster_size`` is extrapolated to this fit's
    realized manifold row count instead of transferred verbatim from the selection run."""
    # Filesystem base path for ``Linker.dump`` (``.gz`` added by the linker).
    model_path: str | None = None
    # Directory for fit-time reports (``linker_fit.clustering_report.json``).
    report_path: str | None = None
    embeddings_parquet: Any = MISSING
    input_text_table_path: str | None = None
    use_gpu: bool = False
    nlp_model: str = "en_core_web_trf"
    # Stage (A): text table I/O buffer rows, encoder batch size (GPU), optional cap on read passes.
    input_buffer_rows: int = 1000
    encoder_batch_size: int = 200
    max_input_buffers: int | None = None
    negatives_per_positive: float = 0.0
    negative_label: str = NEGATIVE_LABEL
    negative_seed: int | None = None
    screener_kind: str = "lda"
    """``lda`` or ``svm``; persisted as :attr:`~pelinker.model.Linker.screener`."""
    projection_enabled: bool = True
    """When false, skip 3D manifold OOV score model (no predict-time gate from that path)."""
    # Stage (B): parquet batching (``batch_size`` rows per read batch).
    batch_size: int = 1000
    kb_name: str | None = None
    kb_version: str = "0.1.0"
    kb_created_at: str | None = None
    kb_description: str = ""
    kb_entity_count: int | None = None
    kb_out_name_min_fraction: float = 0.05
    kb_out_name_top_n: int = 3
    kb_out_ambiguity_min_capture: float = 0.10
    # Discriminator: auto = fit from parquet only if no text table; else embed then fit (legacy).
    # str (not Literal): OmegaConf structured configs reject Literal annotations on fields.
    pipeline: str = "embed_only"
    # Per-parquet backbone/layer (length 1 broadcast, or same length as ``embeddings_parquet``).
    # When omitted, ``model_type`` / ``layers_spec`` scalars apply unless the parquet stem matches
    # ``..._<model>_<layers>`` (see ``_parse_embedding_parquet_stem``).
    model_types: list[str] | None = None
    layers_specs: list[str] | None = None
    # Compact = ParametricUMAP + MLP entity head (default). Legacy = UMAP + HDBSCAN predict.
    predict_mode: str = "compact"
    entity_head_hidden_layers: list[int] | None = None
    """MLP hidden sizes for compact mode; default ``[256, 128, 128]`` when omitted."""
    entity_head_holdout_fraction: float = 0.15
    """Rows withheld from entity-head training to measure distillation fidelity.
    Set ``0.0`` to train on every row and skip the measurement (pre-fidelity behaviour)."""
    entity_head_holdout_group_col: str = "pmid"
    """Column kept whole across the holdout split, so correlated mentions from one
    document cannot straddle it and inflate the measured agreement."""
    entity_head_holdout_seed: int = 13
    distillation_gates_enabled: bool = True
    distillation_min_entity_agreement: float = 0.95
    distillation_max_emit_rate_rel_delta: float = 0.10
    distillation_emit_rate_threshold: float = 0.3
    distillation_on_failure: str = "warn"
    """``warn`` keeps the fitted model and records the failure; ``raise`` aborts the fit."""
    parametric_umap_n_training_epochs: int = 10
    parametric_umap_batch_size: int | None = None

    def __post_init__(self) -> None:
        if self.pipeline not in _PIPELINE_VALUES:
            raise ValueError(
                "pipeline must be one of "
                f"{sorted(_PIPELINE_VALUES)}, got {self.pipeline!r}"
            )
        if self.screener_kind not in ("lda", "svm"):
            raise ValueError(
                f"screener_kind must be 'lda' or 'svm', got {self.screener_kind!r}"
            )
        if self.cluster_viz_method not in ("pca", "umap"):
            raise ValueError(
                f"cluster_viz_method must be 'pca' or 'umap', got {self.cluster_viz_method!r}"
            )
        if self.predict_mode not in ("compact", "legacy"):
            raise ValueError(
                f"predict_mode must be 'compact' or 'legacy', got {self.predict_mode!r}"
            )
        if self.parametric_umap_n_training_epochs < 1:
            raise ValueError("parametric_umap_n_training_epochs must be >= 1")
        if (
            self.parametric_umap_batch_size is not None
            and self.parametric_umap_batch_size < 1
        ):
            raise ValueError("parametric_umap_batch_size must be >= 1 when provided")
        if self.entity_head_hidden_layers is not None:
            if not self.entity_head_hidden_layers or any(
                int(h) < 1 for h in self.entity_head_hidden_layers
            ):
                raise ValueError(
                    "entity_head_hidden_layers must be a non-empty list of ints >= 1"
                )
        if self.min_cluster_size is not None and self.min_cluster_size < 2:
            raise ValueError("min_cluster_size must be >= 2")
        if self.umap_n_neighbors is not None and self.umap_n_neighbors < 2:
            raise ValueError("umap_n_neighbors must be >= 2 when provided")
        if not 0.0 <= self.entity_head_holdout_fraction < 1.0:
            raise ValueError("entity_head_holdout_fraction must be in [0, 1)")
        if self.distillation_on_failure not in ("warn", "raise"):
            raise ValueError(
                "distillation_on_failure must be 'warn' or 'raise', "
                f"got {self.distillation_on_failure!r}"
            )
        if self.clustering_sample_rows is not None and self.clustering_sample_rows < 1:
            raise ValueError("clustering_sample_rows must be >= 1 when provided")
        if self.min_mentions_per_entity < 1:
            raise ValueError("min_mentions_per_entity must be >= 1")
        if (
            self.max_mentions_per_entity is not None
            and self.max_mentions_per_entity < 1
        ):
            raise ValueError("max_mentions_per_entity must be >= 1 when provided")
        if self.max_mentions_negative is not None and self.max_mentions_negative < 1:
            raise ValueError("max_mentions_negative must be >= 1 when provided")
        if self.clustering_sample_index < 0:
            raise ValueError("clustering_sample_index must be >= 0")
        if not 0.0 <= self.kb_out_name_min_fraction <= 1.0:
            raise ValueError("kb_out_name_min_fraction must be in [0, 1]")
        if self.kb_out_name_top_n < 1:
            raise ValueError("kb_out_name_top_n must be >= 1")
        if not 0.0 <= self.kb_out_ambiguity_min_capture <= 1.0:
            raise ValueError("kb_out_ambiguity_min_capture must be in [0, 1]")

clustering_sample_index = 0 class-attribute instance-attribute

Bootstrap index for clustering subsample (match model-selection sample_idx).

clustering_sample_rows = None class-attribute instance-attribute

Max mention rows per clustering bootstrap draw (stratified). None = use all loaded rows.

distillation_on_failure = 'warn' class-attribute instance-attribute

warn keeps the fitted model and records the failure; raise aborts the fit.

entity_head_hidden_layers = None class-attribute instance-attribute

MLP hidden sizes for compact mode; default [256, 128, 128] when omitted.

entity_head_holdout_fraction = 0.15 class-attribute instance-attribute

Rows withheld from entity-head training to measure distillation fidelity. Set 0.0 to train on every row and skip the measurement (pre-fidelity behaviour).

entity_head_holdout_group_col = 'pmid' class-attribute instance-attribute

Column kept whole across the holdout split, so correlated mentions from one document cannot straddle it and inflate the measured agreement.

mention_cap_seed = None class-attribute instance-attribute

Seed for per-entity mention cap; defaults to seed when omitted.

min_cluster_size = None class-attribute instance-attribute

Explicit HDBSCAN min_cluster_size. Omit to resolve from scale_curve_path, or fall back to 20 when neither is given. An explicit value always wins.

pca_components = None class-attribute instance-attribute

PCA components; omit to take the selection report's value, else 100.

pca_seed = 13 class-attribute instance-attribute

Random seed for PCA and cluster-viz PCA.

projection_enabled = True class-attribute instance-attribute

When false, skip 3D manifold OOV score model (no predict-time gate from that path).

scale_curve_path = None class-attribute instance-attribute

scale_curve.json from pelinker-scale-curve. When set (and min_cluster_size is not), min_cluster_size is extrapolated to this fit's realized manifold row count instead of transferred verbatim from the selection run.

screener_kind = 'lda' class-attribute instance-attribute

lda or svm; persisted as :attr:~pelinker.model.Linker.screener.

seed = 13 class-attribute instance-attribute

Bootstrap seed for clustering subsample draws (base_seed); also default for mention-cap and screener draws.

selection_report = None class-attribute instance-attribute

selected_hyperparameters.json (or the report dir containing it) from pelinker-model-selection / pelinker-dim-selection.

Fills in pca_components, umap_dim, umap_n_neighbors and min_cluster_size when those are not set explicitly here. Explicit overrides always win, and any disagreement is logged rather than silently resolved.

umap_dim = None class-attribute instance-attribute

UMAP output dimension; omit to take the selection report's value, else 8.

umap_n_neighbors = None class-attribute instance-attribute

UMAP n_neighbors; omit for the library default (15). Scale-dependent — see pelinker-scale-curve when the fit N differs markedly from the selection N.

umap_seed = None class-attribute instance-attribute

UMAP random seed; omit (None) for parallel UMAP. Set for reproducible production fits.

fit(cfg)

Run embedding (optional), fit a Linker from parquet(s) (optional), and write outputs.

Paths (no implicit fallbacks — missing required paths raise):

• embeddings_parquet: output path(s) for embed_only / both stage (A), or input parquet(s) for fit_only / both stage (B).
• report_path: directory; fit stages write linker_fit.clustering_report.json.gz, linker_fit.cluster_composition.json.gz, and linker_fit.kb_out.json there.
• model_path: filesystem path passed to Linker.dump for fit stages.

Pipelines:

• pipeline=auto: embed then fit if input_text_table_path is set; else fit from parquet.
• pipeline=embed_only: write parquet(s) only (model_path / report_path not used).
• pipeline=fit_only: fit from existing parquet(s); requires model_path and report_path.
• pipeline=both: text table + embed then fit; requires model_path and report_path.

Multiple embeddings_parquet values fuse in list order (inner join on pmid/entity/mention). Set model_types / layers_specs (or scalars) so embedding_metadata.sources matches; or infer model_type / layers_spec from each filename stem when lists are omitted.

Source code in pelinker/cli/fit.py

def fit(cfg: FitCliConfig) -> None:
    """
    Run embedding (optional), fit a ``Linker`` from parquet(s) (optional), and write outputs.

    Paths (no implicit fallbacks — missing required paths raise):

    - ``embeddings_parquet``: output path(s) for ``embed_only`` / ``both`` stage (A), or input
      parquet(s) for ``fit_only`` / ``both`` stage (B).
    - ``report_path``: directory; fit stages write ``linker_fit.clustering_report.json.gz``,
      ``linker_fit.cluster_composition.json.gz``, and ``linker_fit.kb_out.json`` there.
    - ``model_path``: filesystem path passed to ``Linker.dump`` for fit stages.

    Pipelines:

    - ``pipeline=auto``: embed then fit if ``input_text_table_path`` is set; else fit from parquet.
    - ``pipeline=embed_only``: write parquet(s) only (``model_path`` / ``report_path`` not used).
    - ``pipeline=fit_only``: fit from existing parquet(s); requires ``model_path`` and ``report_path``.
    - ``pipeline=both``: text table + embed then fit; requires ``model_path`` and ``report_path``.

    Multiple ``embeddings_parquet`` values fuse in list order (inner join on pmid/entity/mention).
    Set ``model_types`` / ``layers_specs`` (or scalars) so ``embedding_metadata.sources`` matches;
    or infer ``model_type`` / ``layers_spec`` from each filename stem when lists are omitted.
    """
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )

    kb_path, labels_map, _kb_labels = _load_kb_labels_map(cfg)

    resolved = _resolve_selection_hyperparameters(cfg)

    transform_config = TransformConfig(
        pca_components=resolved.pca_components,
        umap_components=resolved.umap_dim,
        umap_n_neighbors=resolved.umap_n_neighbors,
        cluster_viz_method=cfg.cluster_viz_method,
        pca_seed=cfg.pca_seed,
        umap_seed=cfg.umap_seed,
        parametric_umap_n_training_epochs=cfg.parametric_umap_n_training_epochs,
        parametric_umap_batch_size=cfg.parametric_umap_batch_size,
    )

    input_text_table_path = expand_config_path(cfg.input_text_table_path)
    model_path = expand_config_path(cfg.model_path)
    report_path_resolved = expand_config_path(cfg.report_path)

    path_strs = _coerce_str_list(cfg.embeddings_parquet)
    if not path_strs:
        raise ValueError("embeddings_parquet must be one or more paths")

    embed_paths: list[Path] = []
    for s in path_strs:
        p = expand_config_path(s)
        if p is None:
            raise ValueError(f"Invalid embeddings path: {s!r}")
        embed_paths.append(p)

    mts = _coerce_optional_str_list(cfg.model_types)
    lss = _coerce_optional_str_list(cfg.layers_specs)
    embedding_metadata = _embedding_metadata(
        embed_paths, cfg.model_type, cfg.layers_spec, mts, lss
    )

    effective = _resolve_fit_pipeline(
        cfg,
        input_text_table_path=input_text_table_path,
        embed_paths=embed_paths,
        model_path=model_path,
        report_path_resolved=report_path_resolved,
    )

    if effective in ("both", "embed_only"):
        assert input_text_table_path is not None
        _run_embed_stage(
            cfg,
            effective=effective,
            input_text_table_path=input_text_table_path,
            kb_path=kb_path,
            embed_paths=embed_paths,
            embedding_metadata=embedding_metadata,
        )

    if effective == "embed_only":
        logger.info("Embed-only pipeline finished; not fitting or saving a linker.")
        return

    linker_fit_cfg = _build_linker_fit_config(cfg)
    kb_config = _build_kb_config(cfg, kb_path)
    kb_out_naming = KbOutNamingConfig(
        min_fraction=cfg.kb_out_name_min_fraction,
        top_n=cfg.kb_out_name_top_n,
        ambiguity_min_capture=cfg.kb_out_ambiguity_min_capture,
    )

    linker = Linker(
        labels_map=labels_map,
        transform_config=transform_config,
        embedding_metadata=embedding_metadata,
    )

    logger.info("Stage (B): Linker.fit from %s", embed_paths)

    linker.fit(
        embeddings=embed_paths if len(embed_paths) > 1 else embed_paths[0],
        transform_config=transform_config,
        min_cluster_size=resolved.min_cluster_size,
        fit_config=linker_fit_cfg,
        embedding_training=None,
        kb_config=kb_config,
        kb_out_naming=kb_out_naming,
        kb_in_labels_map_path=str(kb_path),
    )

    logger.info("Fitted Linker model with %s KB-out entities", len(linker.vocabulary))
    logger.info(
        "KB-out emergent clusters: %s distinct ids",
        len(linker.cluster_id_to_entity_id),
    )

    if model_path is None or report_path_resolved is None:
        raise ValueError("model_path and report_path must be set when fitting")

    _write_fit_outputs(
        linker,
        cfg,
        model_path=model_path,
        report_path_resolved=report_path_resolved,
    )