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

spacy/default_config.cfg

@@ -95,6 +95,8 @@ 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
+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,216 @@
+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)
+def test_init_nlp_initialize_section_override():
+    """Verify that init_max_examples configured under [initialize] is respected."""
+    train_corpus = CountingCorpus(total_examples=40)
+
+    @spacy.registry.readers("test_counting_train_corpus_init_section")
+    def create_train_corpus(total_examples: int = 40):
+        return train_corpus
+
+    config_str = SAMPLE_LIMIT_CONFIG.replace(
+        "test_counting_train_corpus", "test_counting_train_corpus_init_section"
+    ).replace(
+        "init_max_examples = 5", "init_max_examples = null"
+    ).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 init_max_examples is None 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)