mirror of
https://github.com/explosion/spaCy.git
synced 2024-11-11 12:18:04 +03:00
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import pytest
|
|
from spacy.gold import Example
|
|
|
|
from .util import get_doc
|
|
|
|
from spacy.util import minibatch_by_words
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"doc_sizes, expected_batches",
|
|
[
|
|
([400, 400, 199], [3]),
|
|
([400, 400, 199, 3], [4]),
|
|
([400, 400, 199, 3, 200], [3, 2]),
|
|
|
|
([400, 400, 199, 3, 1], [5]),
|
|
([400, 400, 199, 3, 1, 1500], [5]), # 1500 will be discarded
|
|
([400, 400, 199, 3, 1, 200], [3, 3]),
|
|
([400, 400, 199, 3, 1, 999], [3, 3]),
|
|
([400, 400, 199, 3, 1, 999, 999], [3, 2, 1, 1]),
|
|
|
|
([1, 2, 999], [3]),
|
|
([1, 2, 999, 1], [4]),
|
|
([1, 200, 999, 1], [2, 2]),
|
|
([1, 999, 200, 1], [2, 2]),
|
|
],
|
|
)
|
|
def test_util_minibatch(doc_sizes, expected_batches):
|
|
docs = [get_doc(doc_size) for doc_size in doc_sizes]
|
|
examples = [Example(doc=doc) for doc in docs]
|
|
tol = 0.2
|
|
batch_size = 1000
|
|
batches = list(minibatch_by_words(examples=examples, size=batch_size, tolerance=tol, discard_oversize=True))
|
|
assert [len(batch) for batch in batches] == expected_batches
|
|
|
|
max_size = batch_size + batch_size * tol
|
|
for batch in batches:
|
|
assert sum([len(example.doc) for example in batch]) < max_size
|
|
|