fix to pretrain script (#5699)

* fix to pretrain script

* remove unnecessary import
This commit is contained in:
Sofie Van Landeghem 2020-07-02 21:48:01 +02:00 committed by GitHub
parent a723fa02a1
commit 41b65fd0f8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -15,7 +15,6 @@ from ..ml.models.multi_task import build_masked_language_model
from ..tokens import Doc
from ..attrs import ID, HEAD
from .. import util
from ..gold import Example
@app.command("pretrain")
@ -183,7 +182,7 @@ def pretrain(
for batch_id, batch in enumerate(batches):
docs, count = make_docs(
nlp,
[ex.doc for ex in batch],
batch,
max_length=pretrain_config["max_length"],
min_length=pretrain_config["min_length"],
)