Skip to content

Model setup

gutenberg_kg.model_setup

model_setup.py - Ensure local ML models used by the pipeline are present.

Two models are needed on the machine that runs chunk-diaries / ingest / build-corpus (they are never downloaded at Docker runtime — see docker/Dockerfile, which pre-downloads only the embedder and never touches spaCy since chunking happens locally before the image is built):

  • en_core_web_sm spaCy model used by diary-transformer for sentence segmentation (gutenkg chunk-diaries).
  • BAAI/bge-small-en-v1.5 sentence-transformers embedder used by both gutenkg ingest and gutenkg build-corpus (must match EMBED_MODEL in the Docker image).

Neither is a pip dependency: the spaCy model ships via spacy download, and the embedder is fetched from HuggingFace on first use. A fresh clone + venv has neither, and the failure only surfaces mid-pipeline (see diary/chunk.py), which is what this module lets gutenkg init catch up front instead.

ModelCheckResult(name, kind, status, message='') dataclass

Outcome of checking (and optionally fetching) one model.

check_models(download=True)

Check (and by default fetch) every model the local pipeline needs.

Parameters:

Name Type Description Default
download bool

When True, missing models are downloaded. When False, this only reports what's missing (use for a dry --check).

True

Returns:

Type Description
list[ModelCheckResult]

One :class:ModelCheckResult per model, spaCy first.

Source code in src/gutenberg_kg/model_setup.py
92
93
94
95
96
97
98
99
def check_models(download: bool = True) -> list[ModelCheckResult]:
    """Check (and by default fetch) every model the local pipeline needs.

    :param download: When True, missing models are downloaded. When False,
        this only reports what's missing (use for a dry ``--check``).
    :returns: One :class:`ModelCheckResult` per model, spaCy first.
    """
    return [_check_spacy(download), _check_embedder(download)]