feat(training): add init_max_examples to limit train corpus during initialization (#13804)

spacy/default_config.cfg

@@ -95,6 +95,9 @@ train_corpus = "corpora.train"
 before_to_disk = null
 # Optional callback that is invoked at the start of each training step
 before_update = null
+# Maximum number of training examples to sample for pipeline initialization.
+# null or 0 means unlimited (or 100 if max_epochs = -1).
+init_max_examples = null
 
 [training.logger]
 @loggers = "spacy.ConsoleLogger.v1"

spacy/schemas.py

@@ -372,6 +372,7 @@ class ConfigSchemaTraining(BaseModel):
     annotating_components: List[str] = Field(..., title="Pipeline components that should set annotations during training")
     before_to_disk: Optional[Callable[["Language"], "Language"]] = Field(..., title="Optional callback to modify nlp object after training, before it's saved to disk")
     before_update: Optional[Callable[["Language", Dict[str, Any]], None]] = Field(..., title="Optional callback that is invoked at the start of each training step")
+    init_max_examples: Optional[StrictInt] = Field(None, ge=0, title="Maximum number of training examples to sample for pipeline initialization")
     # fmt: on
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
@@ -424,6 +425,7 @@ class ConfigSchemaInit(BaseModel):
     components: Dict[StrictStr, Dict[StrictStr, Any]] = Field(..., title="Arguments for TrainablePipe.initialize methods of pipeline components, keyed by component")
     before_init: Optional[Callable[["Language"], "Language"]] = Field(..., title="Optional callback to modify nlp object before initialization")
     after_init: Optional[Callable[["Language"], "Language"]] = Field(..., title="Optional callback to modify nlp object after initialization")
+    init_max_examples: Optional[StrictInt] = Field(None, ge=0, title="Maximum number of training examples to sample for pipeline initialization")
     # fmt: on
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

spacy/tests/training/test_initialize_sample_limit.py

@@ -0,0 +1,218 @@
+import pytest
+from typing import Iterable
+import spacy
+from spacy.training import Example
+from spacy.training.initialize import init_nlp
+from spacy import util
+from thinc.api import ConfigValidationError
+
+
+class CountingCorpus:
+    """Mock corpus reader that records how many examples were yielded per generator call."""
+
+    def __init__(self, total_examples: int = 50):
+        self.total_examples = total_examples
+        self.call_counts = []
+
+    def __call__(self, nlp) -> Iterable[Example]:
+        nlp_blank = spacy.blank("en")
+        idx = len(self.call_counts)
+        self.call_counts.append(0)
+        for i in range(self.total_examples):
+            self.call_counts[idx] += 1
+            doc = nlp_blank.make_doc(f"This is example number {i}")
+            doc.cats = {"LABEL": 1.0}
+            yield Example.from_dict(doc, {"cats": {"LABEL": 1.0}})
+
+
+SAMPLE_LIMIT_CONFIG = """
+[paths]
+train = null
+dev = null
+vectors = null
+init_tok2vec = null
+
+[system]
+gpu_allocator = null
+seed = 0
+
+[nlp]
+lang = "en"
+pipeline = ["textcat_multilabel"]
+
+[components]
+
+[components.textcat_multilabel]
+factory = "textcat_multilabel"
+
+[corpora]
+
+[corpora.train]
+@readers = "test_counting_train_corpus"
+total_examples = 50
+
+[corpora.dev]
+@readers = "test_counting_dev_corpus"
+
+[training]
+train_corpus = "corpora.train"
+dev_corpus = "corpora.dev"
+seed = ${system.seed}
+gpu_allocator = ${system.gpu_allocator}
+dropout = 0.1
+accumulate_gradient = 1
+patience = 1600
+max_epochs = 5
+max_steps = 1000
+eval_frequency = 100
+score_weights = {}
+frozen_components = []
+annotating_components = []
+before_to_disk = null
+before_update = null
+init_max_examples = 5
+
+[training.batcher]
+@batchers = "spacy.batch_by_words.v1"
+discard_oversize = false
+tolerance = 0.2
+get_length = null
+
+[training.batcher.size]
+@schedules = "compounding.v1"
+start = 100
+stop = 1000
+compound = 1.001
+t = 0.0
+
+[training.optimizer]
+@optimizers = "Adam.v1"
+beta1 = 0.9
+beta2 = 0.999
+L2_is_weight_decay = true
+L2 = 0.01
+grad_clip = 1.0
+use_averages = false
+eps = 0.00000001
+learn_rate = 0.001
+
+[training.logger]
+@loggers = "spacy.ConsoleLogger.v1"
+progress_bar = false
+
+[initialize]
+vectors = ${paths.vectors}
+init_tok2vec = ${paths.init_tok2vec}
+vocab_data = null
+lookups = null
+tokenizer = {}
+before_init = null
+after_init = null
+
+[initialize.components]
+
+[initialize.components.textcat_multilabel]
+labels = ["LABEL"]
+"""
+
+
+@pytest.mark.issue(13804)
+def test_init_nlp_corpus_sample_limit():
+    """Verify that init_max_examples bounds each generator iteration during init_nlp."""
+    train_corpus = CountingCorpus(total_examples=50)
+
+    @spacy.registry.readers("test_counting_train_corpus")
+    def create_train_corpus(total_examples: int = 50):
+        return train_corpus
+
+    @spacy.registry.readers("test_counting_dev_corpus")
+    def create_dev_corpus():
+        return CountingCorpus(total_examples=10)
+
+    config = util.load_config_from_str(SAMPLE_LIMIT_CONFIG)
+    init_nlp(config)
+
+    # When init_max_examples = 5, no generator pass should yield more than 5 examples
+    assert len(train_corpus.call_counts) > 0
+    assert max(train_corpus.call_counts) == 5
+
+
+@pytest.mark.issue(13804)
+def test_init_nlp_default_unbounded_sample():
+    """Verify that when init_max_examples is None, the full finite corpus is evaluated."""
+    train_corpus = CountingCorpus(total_examples=30)
+
+    @spacy.registry.readers("test_counting_train_corpus_unbounded")
+    def create_train_corpus_unbounded(total_examples: int = 30):
+        return train_corpus
+
+    config_str = SAMPLE_LIMIT_CONFIG.replace(
+        "test_counting_train_corpus", "test_counting_train_corpus_unbounded"
+    ).replace("init_max_examples = 5", "init_max_examples = null")
+    config = util.load_config_from_str(config_str)
+    init_nlp(config)
+
+    # All 30 examples should have been yielded in at least one pass
+    assert len(train_corpus.call_counts) > 0
+    assert max(train_corpus.call_counts) == 30
+
+
+@pytest.mark.issue(13804)
+@pytest.mark.parametrize("training_val", ["null", "0"])
+def test_init_nlp_initialize_section_override(training_val):
+    """Verify that init_max_examples under [initialize] is respected when [training] is null or 0."""
+    train_corpus = CountingCorpus(total_examples=40)
+
+    @spacy.registry.readers(f"test_counting_train_corpus_init_section_{training_val}")
+    def create_train_corpus(total_examples: int = 40):
+        return train_corpus
+
+    config_str = SAMPLE_LIMIT_CONFIG.replace(
+        "test_counting_train_corpus",
+        f"test_counting_train_corpus_init_section_{training_val}",
+    ).replace(
+        "init_max_examples = 5", f"init_max_examples = {training_val}"
+    ).replace(
+        "[initialize]\nvectors = ${paths.vectors}",
+        "[initialize]\ninit_max_examples = 7\nvectors = ${paths.vectors}",
+    )
+    config = util.load_config_from_str(config_str)
+    init_nlp(config)
+
+    assert len(train_corpus.call_counts) > 0
+    assert max(train_corpus.call_counts) == 7
+
+
+@pytest.mark.issue(13804)
+def test_init_nlp_streamed_corpus_custom_limit():
+    """Verify that init_max_examples overrides the 100 default when max_epochs = -1."""
+    train_corpus = CountingCorpus(total_examples=50)
+
+    @spacy.registry.readers("test_counting_train_corpus_streamed")
+    def create_train_corpus(total_examples: int = 50):
+        return train_corpus
+
+    config_str = SAMPLE_LIMIT_CONFIG.replace(
+        "test_counting_train_corpus", "test_counting_train_corpus_streamed"
+    ).replace(
+        "max_epochs = 5", "max_epochs = -1"
+    ).replace(
+        "init_max_examples = 5", "init_max_examples = 12"
+    )
+    config = util.load_config_from_str(config_str)
+    init_nlp(config)
+
+    assert len(train_corpus.call_counts) > 0
+    assert max(train_corpus.call_counts) == 12
+
+
+@pytest.mark.issue(13804)
+def test_init_nlp_negative_limit_rejected():
+    """Verify that negative init_max_examples is rejected by schema validation."""
+    config_str = SAMPLE_LIMIT_CONFIG.replace(
+        "init_max_examples = 5", "init_max_examples = -5"
+    )
+    config = util.load_config_from_str(config_str)
+    with pytest.raises(ConfigValidationError):
+        init_nlp(config)
+

spacy/training/initialize.py

@@ -79,8 +79,15 @@ def init_nlp(config: Config, *, use_gpu: int = -1) -> "Language":
     # defined before initializing further
     nlp._link_components()
     with nlp.select_pipes(disable=[*frozen_components, *resume_components]):
+        init_max_examples = T.get("init_max_examples")
+        if not init_max_examples and "initialize" in config:
+            init_max_examples = config.get("initialize", {}).get("init_max_examples")
         if T["max_epochs"] == -1:
-            sample_size = 100
+            sample_size = (
+                init_max_examples
+                if init_max_examples is not None and init_max_examples > 0
+                else 100
+            )
             logger.debug(
                 "Due to streamed train corpus, using only first %s examples for initialization. "
                 "If necessary, provide all labels in [initialize]. "
@@ -90,6 +97,14 @@ def init_nlp(config: Config, *, use_gpu: int = -1) -> "Language":
             nlp.initialize(
                 lambda: islice(train_corpus(nlp), sample_size), sgd=optimizer
             )
+        elif init_max_examples is not None and init_max_examples > 0:
+            logger.info(
+                "Using first %s examples of train corpus for initialization.",
+                init_max_examples,
+            )
+            nlp.initialize(
+                lambda: islice(train_corpus(nlp), init_max_examples), sgd=optimizer
+            )
         else:
             nlp.initialize(lambda: train_corpus(nlp), sgd=optimizer)
         logger.info("Initialized pipeline components: %s", nlp.pipe_names)

website/docs/api/data-formats.mdx

@@ -193,6 +193,7 @@ process that are used when you run [`spacy train`](/api/cli#train).
 | `frozen_components`                                  | Pipeline component names that are "frozen" and shouldn't be initialized or updated during training. See [here](/usage/training#config-components) for details. Defaults to `[]`. ~~List[str]~~                                                                                                                                      |
 | `annotating_components` <Tag variant="new">3.1</Tag> | Pipeline component names that should set annotations on the predicted docs during training. See [here](/usage/training#annotating-components) for details. Defaults to `[]`. ~~List[str]~~                                                                                                                                          |
 | `gpu_allocator`                                      | Library for cupy to route GPU memory allocation to. Can be `"pytorch"` or `"tensorflow"`. Defaults to variable `${system.gpu_allocator}`. ~~str~~                                                                                                                                                                                   |
+| `init_max_examples`                                  | Maximum number of training examples to sample for pipeline initialization. `0` or `null` means an unlimited number of examples (or `100` if `max_epochs = -1`). Defaults to `null`. ~~Optional[int]~~                                                                                                                        |
 | `logger`                                             | Callable that takes the `nlp` and stdout and stderr `IO` objects, sets up the logger, and returns two new callables to log a training step and to finalize the logger. Defaults to [`ConsoleLogger`](/api/top-level#ConsoleLogger). ~~Callable[[Language, IO, IO], [Tuple[Callable[[Dict[str, Any]], None], Callable[[], None]]]]~~ |
 | `max_epochs`                                         | Maximum number of epochs to train for. `0` means an unlimited number of epochs. `-1` means that the train corpus should be streamed rather than loaded into memory with no shuffling within the training loop. Defaults to `0`. ~~int~~                                                                                             |
 | `max_steps`                                          | Maximum number of update steps to train for. `0` means an unlimited number of steps. Defaults to `20000`. ~~int~~                                                                                                                                                                                                                   |
@@ -248,8 +249,9 @@ Also see the usage guides on the
 | -------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | `after_init`   | Optional callback to modify the `nlp` object after initialization. ~~Optional[Callable[[Language], Language]]~~                                                                                                                                                                                                                                                                                                |
 | `before_init`  | Optional callback to modify the `nlp` object before initialization. ~~Optional[Callable[[Language], Language]]~~                                                                                                                                                                                                                                                                                               |
-| `components`   | Additional arguments passed to the `initialize` method of a pipeline component, keyed by component name. If type annotations are available on the method, the config will be validated against them. The `initialize` methods will always receive the `get_examples` callback and the current `nlp` object. ~~Dict[str, Dict[str, Any]]~~                                                                      |
-| `init_tok2vec` | Optional path to pretrained tok2vec weights created with [`spacy pretrain`](/api/cli#pretrain). Defaults to variable `${paths.init_tok2vec}`. Ignored when actually running pretraining, as you're creating the file to be used later. ~~Optional[str]~~                                                                                                                                                       |
+| `components`        | Additional arguments passed to the `initialize` method of a pipeline component, keyed by component name. If type annotations are available on the method, the config will be validated against them. The `initialize` methods will always receive the `get_examples` callback and the current `nlp` object. ~~Dict[str, Dict[str, Any]]~~                                                                      |
+| `init_max_examples` | Optional maximum number of training examples to sample for pipeline initialization. Overrides or provides fallback for `[training.init_max_examples]`. `0` or `null` means unlimited. Defaults to `null`. ~~Optional[int]~~                                                                                                                                                                                      |
+| `init_tok2vec`      | Optional path to pretrained tok2vec weights created with [`spacy pretrain`](/api/cli#pretrain). Defaults to variable `${paths.init_tok2vec}`. Ignored when actually running pretraining, as you're creating the file to be used later. ~~Optional[str]~~                                                                                                                                                       |
 | `lookups`      | Additional lexeme and vocab data from [`spacy-lookups-data`](https://github.com/explosion/spacy-lookups-data). Defaults to `null`. ~~Optional[Lookups]~~                                                                                                                                                                                                                                                       |
 | `tokenizer`    | Additional arguments passed to the `initialize` method of the specified tokenizer. Can be used for languages like Chinese that depend on dictionaries or trained models for tokenization. If type annotations are available on the method, the config will be validated against them. The `initialize` method will always receive the `get_examples` callback and the current `nlp` object. ~~Dict[str, Any]~~ |
 | `vectors`      | Name or path of pipeline containing pretrained word vectors to use, e.g. created with [`init vectors`](/api/cli#init-vectors). Defaults to `null`. ~~Optional[str]~~                                                                                                                                                                                                                                           |