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

[JohnnyWilson16]

Description

Fixes #13804.

During spacy train (and spacy.training.initialize.init_nlp), pipeline components are initialized by calling nlp.initialize() with an unbounded generator lambda: train_corpus(nlp) for non-streaming finite datasets. On large datasets (e.g. 10GB+ DocBin with millions of documents), components such as textcat, tagger, ner iterate through the entire dataset just to extract labels and configure dimensions, which can stall the CLI for hours before training starts.

Following guidance from @honnibal on #13804:

• Adds configurable init_max_examples: Optional[StrictInt] = Field(None, ge=0, ...) to [training] and [initialize] (ConfigSchemaTraining & ConfigSchemaInit).
• When init_max_examples is set (e.g. init_max_examples = 2000), train_corpus(nlp) is bounded during pipeline initialization using itertools.islice(train_corpus(nlp), init_max_examples).
• When streaming corpora are used (max_epochs = -1), init_max_examples can be used to customize the default 100-sample limit.
• Retains existing unconstrained behavior when init_max_examples is null/None or 0 for full backward compatibility with existing configs.
• Rejects negative integers via Pydantic/confection schema validation.
• Unit and regression tests added under spacy/tests/training/test_initialize_sample_limit.py.

Types of change

• Enhancement / new feature

Checklist

• I confirm that I have the right to submit this contribution under the project's MIT license.
• I ran the tests, and all new and existing tests passed.
• My changes don't require a change to the documentation, or if they do, I've added all required information.

[chrikrah]

@JohnnyWilson16 house style says 0 is unbounded, as in max_epochs, max_steps and patience just above it. Your fourth bullet is not what I question. Precedence is. In initialize.py, 0 from [training] is set for the fallback test and unset for the limit test.

blocking: [training] init_max_examples = 0 swallows [initialize] init_max_examples = 7, the case test_init_nlp_initialize_section_override exists to cover.

$ cat probe.py          # 088d6990, spaCy 3.8.16 built from this tree, CPython 3.12.3
import spacy
from spacy import util
from spacy.training.initialize import init_nlp
from spacy.tests.training import test_initialize_sample_limit as t

def pulled(training_value):
    corpus = t.CountingCorpus(total_examples=50)
    spacy.registry.readers("pp_train")(lambda total_examples=50: corpus)
    spacy.registry.readers("pp_dev")(lambda: t.CountingCorpus(total_examples=10))
    cfg = t.SAMPLE_LIMIT_CONFIG
    cfg = cfg.replace('"test_counting_train_corpus"', '"pp_train"')
    cfg = cfg.replace('"test_counting_dev_corpus"', '"pp_dev"')
    cfg = cfg.replace("init_max_examples = 5", "init_max_examples = " + training_value)
    cfg = cfg.replace("[initialize]\n", "[initialize]\ninit_max_examples = 7\n")
    init_nlp(util.load_config_from_str(cfg))
    return max(corpus.call_counts)

for value in ["null", "0"]:
    print("[training] %-4s + [initialize] 7 -> %d examples pulled" % (value, pulled(value)))

$ python probe.py
[training] null + [initialize] 7 -> 7 examples pulled
[training] 0    + [initialize] 7 -> 50 examples pulled

One line at initialize.py:83 sends 0 to the [initialize] lookup. It stays unbounded when nothing is there, which your description already promises.

$ sed -i 's/if init_max_examples is None and/if not init_max_examples and/' spacy/training/initialize.py
$ python probe.py       # 088d6990 plus that one line
[training] null + [initialize] 7 -> 7 examples pulled
[training] 0    + [initialize] 7 -> 7 examples pulled
$ python -m pytest spacy/tests/training/test_initialize_sample_limit.py -q -p no:randomly
5 passed in 1.25s       # 4 failed, 1 passed with the c2dabfce source files; not run: spacy/tests/pipeline

non-blocking: the key has no row in the [training] or [initialize] tables of website/docs/api/data-formats.mdx, where every sibling has one. Its default_config.cfg comment also omits what the value does when unset.

Does the one-liner suit you, or do you want [training] 0 to keep beating [initialize]?

[JohnnyWilson16]

@chrikrah The precedence fix and documentation updates have been pushed in fb9b8b9:

• Updated spacy/training/initialize.py to use if not init_max_examples and "initialize" in config: so that 0 from [training] defers to [initialize].
• Parameterized test_init_nlp_initialize_section_override across both null and 0 to verify this behavior.
• Added init_max_examples rows to the [training] and [initialize] tables in website/docs/api/data-formats.mdx.
• Clarified in spacy/default_config.cfg that null or 0 means unlimited (or 100 if max_epochs = -1).

[chrikrah]

@JohnnyWilson16 approving at fb9b8b9. [training] init_max_examples = 0 now defers to [initialize], and the parameterised test fails without that line.

$ python probe.py          # fb9b8b9, spaCy 3.8.16 built from this tree, CPython 3.12.3
[training] null + [initialize] 7    -> 7 examples pulled
[training] 0    + [initialize] 7    -> 7 examples pulled
[training] 5    + [initialize] 7    -> 5 examples pulled
[training] 0    + [initialize] null -> 50 examples pulled

$ python -m pytest spacy/tests/training/test_initialize_sample_limit.py -q -p no:randomly
6 passed                   # fb9b8b9
1 failed, 5 passed         # initialize.py taken from 088d6990: test_init_nlp_initialize_section_override[0]

$ python -m pytest spacy/tests/training -q -p no:randomly
117 passed, 3 skipped      # fb9b8b9
# not run: spacy/tests/pipeline, and no website build

probe.py, the one from my first review with a second argument

import spacy
from spacy import util
from spacy.training.initialize import init_nlp
from spacy.tests.training import test_initialize_sample_limit as t

def pulled(training_value, initialize_value):
    corpus = t.CountingCorpus(total_examples=50)
    name = "pp_train_%s_%s" % (training_value, initialize_value)
    spacy.registry.readers(name)(lambda total_examples=50: corpus)
    spacy.registry.readers("pp_dev")(lambda: t.CountingCorpus(total_examples=10))
    cfg = t.SAMPLE_LIMIT_CONFIG
    cfg = cfg.replace('"test_counting_train_corpus"', '"%s"' % name)
    cfg = cfg.replace('"test_counting_dev_corpus"', '"pp_dev"')
    cfg = cfg.replace("init_max_examples = 5", "init_max_examples = " + training_value)
    cfg = cfg.replace("[initialize]\n", "[initialize]\ninit_max_examples = %s\n" % initialize_value)
    init_nlp(util.load_config_from_str(cfg))
    return max(corpus.call_counts)

for tv, iv in [("null", "7"), ("0", "7"), ("5", "7"), ("0", "null")]:
    print("[training] %-4s + [initialize] %-4s -> %d examples pulled" % (tv, iv, pulled(tv, iv)))

non-blocking: the new [initialize] row in data-formats.mdx says the key "Overrides or provides fallback for" the [training] one. The third probe line shows [training] 5 winning over [initialize] 7, so "Used when [training] init_max_examples is null or 0" is what the code does.

@honnibal, would you take a look at this one?