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
    pca_components: int = 100
    umap_dim: int = 8
    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 = 20
    # 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
    # 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

    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.min_cluster_size < 2:
            raise ValueError("min_cluster_size must be >= 2")
        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")

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.

mention_cap_seed = None class-attribute instance-attribute

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

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).

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.

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 and linker_fit.cluster_composition.json.gz there (see :func:pelinker.reporting.linker_fit_clustering_report_path and :func:pelinker.reporting.linker_fit_cluster_composition_path).
• 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`` and
      ``linker_fit.cluster_composition.json.gz`` there (see
      :func:`pelinker.reporting.linker_fit_clustering_report_path` and
      :func:`pelinker.reporting.linker_fit_cluster_composition_path`).
    - ``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)

    transform_config = TransformConfig(
        pca_components=cfg.pca_components,
        umap_components=cfg.umap_dim,
        cluster_viz_method=cfg.cluster_viz_method,
        pca_seed=cfg.pca_seed,
        umap_seed=cfg.umap_seed,
    )

    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)

    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=cfg.min_cluster_size,
        fit_config=linker_fit_cfg,
        embedding_training=None,
        kb_config=kb_config,
    )

    logger.info("Fitted Linker model with %s entities", len(linker.vocabulary))
    logger.info(
        "Entity-level provisional clusters: %s distinct ids",
        len(set(linker.cluster_assignments.values())),
    )

    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,
    )