We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Pretraining: batch size & document length
## Feature description

I was just testing the `pretrain` command and ran out of memory ^^ !
Looking at the code, it seems that the `batch_size` is fixed to 3000. I also noticed that document with length greater than 500 token are pruned by default. I guess this are default parameters for a corpus of sentence / small text... But maybe we should let the user set these parameters ?



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `spacy/cli/pretrain.py`
Content:
```
1 # coding: utf8
2 from __future__ import print_function, unicode_literals
3 
4 import plac
5 import random
6 import numpy
7 import time
8 from collections import Counter
9 from pathlib import Path
10 from thinc.v2v import Affine, Maxout
11 from thinc.misc import LayerNorm as LN
12 from thinc.neural.util import prefer_gpu
13 from wasabi import Printer
14 import srsly
15 
16 from ..tokens import Doc
17 from ..attrs import ID, HEAD
18 from .._ml import Tok2Vec, flatten, chain, create_default_optimizer
19 from .._ml import masked_language_model
20 from .. import util
21 
22 
23 @plac.annotations(
24     texts_loc=("Path to jsonl file with texts to learn from", "positional", None, str),
25     vectors_model=("Name or path to vectors model to learn from"),
26     output_dir=("Directory to write models each epoch", "positional", None, str),
27     width=("Width of CNN layers", "option", "cw", int),
28     depth=("Depth of CNN layers", "option", "cd", int),
29     embed_rows=("Embedding rows", "option", "er", int),
30     use_vectors=("Whether to use the static vectors as input features", "flag", "uv"),
31     dropout=("Dropout", "option", "d", float),
32     seed=("Seed for random number generators", "option", "s", float),
33     nr_iter=("Number of iterations to pretrain", "option", "i", int),
34 )
35 def pretrain(
36     texts_loc,
37     vectors_model,
38     output_dir,
39     width=96,
40     depth=4,
41     embed_rows=2000,
42     use_vectors=False,
43     dropout=0.2,
44     nr_iter=1000,
45     seed=0,
46 ):
47     """
48     Pre-train the 'token-to-vector' (tok2vec) layer of pipeline components,
49     using an approximate language-modelling objective. Specifically, we load
50     pre-trained vectors, and train a component like a CNN, BiLSTM, etc to predict
51     vectors which match the pre-trained ones. The weights are saved to a directory
52     after each epoch. You can then pass a path to one of these pre-trained weights
53     files to the 'spacy train' command.
54 
55     This technique may be especially helpful if you have little labelled data.
56     However, it's still quite experimental, so your mileage may vary.
57 
58     To load the weights back in during 'spacy train', you need to ensure
59     all settings are the same between pretraining and training. The API and
60     errors around this need some improvement.
61     """
62     config = dict(locals())
63     msg = Printer()
64     util.fix_random_seed(seed)
65 
66     has_gpu = prefer_gpu()
67     msg.info("Using GPU" if has_gpu else "Not using GPU")
68 
69     output_dir = Path(output_dir)
70     if not output_dir.exists():
71         output_dir.mkdir()
72         msg.good("Created output directory")
73     srsly.write_json(output_dir / "config.json", config)
74     msg.good("Saved settings to config.json")
75 
76     # Load texts from file or stdin
77     if texts_loc != "-":  # reading from a file
78         texts_loc = Path(texts_loc)
79         if not texts_loc.exists():
80             msg.fail("Input text file doesn't exist", texts_loc, exits=1)
81         with msg.loading("Loading input texts..."):
82             texts = list(srsly.read_jsonl(texts_loc))
83         msg.good("Loaded input texts")
84         random.shuffle(texts)
85     else:  # reading from stdin
86         msg.text("Reading input text from stdin...")
87         texts = srsly.read_jsonl("-")
88 
89     with msg.loading("Loading model '{}'...".format(vectors_model)):
90         nlp = util.load_model(vectors_model)
91     msg.good("Loaded model '{}'".format(vectors_model))
92     pretrained_vectors = None if not use_vectors else nlp.vocab.vectors.name
93     model = create_pretraining_model(
94         nlp,
95         Tok2Vec(
96             width,
97             embed_rows,
98             conv_depth=depth,
99             pretrained_vectors=pretrained_vectors,
100             bilstm_depth=0,  # Requires PyTorch. Experimental.
101             cnn_maxout_pieces=3,  # You can try setting this higher
102             subword_features=True,  # Set to False for Chinese etc
103         ),
104     )
105     optimizer = create_default_optimizer(model.ops)
106     tracker = ProgressTracker(frequency=10000)
107     msg.divider("Pre-training tok2vec layer")
108     row_settings = {"widths": (3, 10, 10, 6, 4), "aligns": ("r", "r", "r", "r", "r")}
109     msg.row(("#", "# Words", "Total Loss", "Loss", "w/s"), **row_settings)
110     for epoch in range(nr_iter):
111         for batch in util.minibatch_by_words(
112             ((text, None) for text in texts), size=3000
113         ):
114             docs = make_docs(nlp, [text for (text, _) in batch])
115             loss = make_update(model, docs, optimizer, drop=dropout)
116             progress = tracker.update(epoch, loss, docs)
117             if progress:
118                 msg.row(progress, **row_settings)
119                 if texts_loc == "-" and tracker.words_per_epoch[epoch] >= 10 ** 7:
120                     break
121         with model.use_params(optimizer.averages):
122             with (output_dir / ("model%d.bin" % epoch)).open("wb") as file_:
123                 file_.write(model.tok2vec.to_bytes())
124             log = {
125                 "nr_word": tracker.nr_word,
126                 "loss": tracker.loss,
127                 "epoch_loss": tracker.epoch_loss,
128                 "epoch": epoch,
129             }
130             with (output_dir / "log.jsonl").open("a") as file_:
131                 file_.write(srsly.json_dumps(log) + "\n")
132         tracker.epoch_loss = 0.0
133         if texts_loc != "-":
134             # Reshuffle the texts if texts were loaded from a file
135             random.shuffle(texts)
136 
137 
138 def make_update(model, docs, optimizer, drop=0.0, objective="L2"):
139     """Perform an update over a single batch of documents.
140 
141     docs (iterable): A batch of `Doc` objects.
142     drop (float): The droput rate.
143     optimizer (callable): An optimizer.
144     RETURNS loss: A float for the loss.
145     """
146     predictions, backprop = model.begin_update(docs, drop=drop)
147     loss, gradients = get_vectors_loss(model.ops, docs, predictions, objective)
148     backprop(gradients, sgd=optimizer)
149     # Don't want to return a cupy object here
150     # The gradients are modified in-place by the BERT MLM,
151     # so we get an accurate loss
152     return float(loss)
153 
154 
155 def make_docs(nlp, batch, min_length=1, max_length=500):
156     docs = []
157     for record in batch:
158         text = record["text"]
159         if "tokens" in record:
160             doc = Doc(nlp.vocab, words=record["tokens"])
161         else:
162             doc = nlp.make_doc(text)
163         if "heads" in record:
164             heads = record["heads"]
165             heads = numpy.asarray(heads, dtype="uint64")
166             heads = heads.reshape((len(doc), 1))
167             doc = doc.from_array([HEAD], heads)
168         if len(doc) >= min_length and len(doc) < max_length:
169             docs.append(doc)
170     return docs
171 
172 
173 def get_vectors_loss(ops, docs, prediction, objective="L2"):
174     """Compute a mean-squared error loss between the documents' vectors and
175     the prediction.
176 
177     Note that this is ripe for customization! We could compute the vectors
178     in some other word, e.g. with an LSTM language model, or use some other
179     type of objective.
180     """
181     # The simplest way to implement this would be to vstack the
182     # token.vector values, but that's a bit inefficient, especially on GPU.
183     # Instead we fetch the index into the vectors table for each of our tokens,
184     # and look them up all at once. This prevents data copying.
185     ids = ops.flatten([doc.to_array(ID).ravel() for doc in docs])
186     target = docs[0].vocab.vectors.data[ids]
187     if objective == "L2":
188         d_scores = prediction - target
189         loss = (d_scores ** 2).sum()
190     else:
191         raise NotImplementedError(objective)
192     return loss, d_scores
193 
194 
195 def create_pretraining_model(nlp, tok2vec):
196     """Define a network for the pretraining. We simply add an output layer onto
197     the tok2vec input model. The tok2vec input model needs to be a model that
198     takes a batch of Doc objects (as a list), and returns a list of arrays.
199     Each array in the output needs to have one row per token in the doc.
200     """
201     output_size = nlp.vocab.vectors.data.shape[1]
202     output_layer = chain(
203         LN(Maxout(300, pieces=3)), Affine(output_size, drop_factor=0.0)
204     )
205     # This is annoying, but the parser etc have the flatten step after
206     # the tok2vec. To load the weights in cleanly, we need to match
207     # the shape of the models' components exactly. So what we cann
208     # "tok2vec" has to be the same set of processes as what the components do.
209     tok2vec = chain(tok2vec, flatten)
210     model = chain(tok2vec, output_layer)
211     model = masked_language_model(nlp.vocab, model)
212     model.tok2vec = tok2vec
213     model.output_layer = output_layer
214     model.begin_training([nlp.make_doc("Give it a doc to infer shapes")])
215     return model
216 
217 
218 class ProgressTracker(object):
219     def __init__(self, frequency=1000000):
220         self.loss = 0.0
221         self.prev_loss = 0.0
222         self.nr_word = 0
223         self.words_per_epoch = Counter()
224         self.frequency = frequency
225         self.last_time = time.time()
226         self.last_update = 0
227         self.epoch_loss = 0.0
228 
229     def update(self, epoch, loss, docs):
230         self.loss += loss
231         self.epoch_loss += loss
232         words_in_batch = sum(len(doc) for doc in docs)
233         self.words_per_epoch[epoch] += words_in_batch
234         self.nr_word += words_in_batch
235         words_since_update = self.nr_word - self.last_update
236         if words_since_update >= self.frequency:
237             wps = words_since_update / (time.time() - self.last_time)
238             self.last_update = self.nr_word
239             self.last_time = time.time()
240             loss_per_word = self.loss - self.prev_loss
241             status = (
242                 epoch,
243                 self.nr_word,
244                 "%.8f" % self.loss,
245                 "%.8f" % loss_per_word,
246                 int(wps),
247             )
248             self.prev_loss = float(self.loss)
249             return status
250         else:
251             return None
252
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/spacy/cli/pretrain.py b/spacy/cli/pretrain.py
--- a/spacy/cli/pretrain.py
+++ b/spacy/cli/pretrain.py
@@ -29,6 +29,9 @@
     embed_rows=("Embedding rows", "option", "er", int),
     use_vectors=("Whether to use the static vectors as input features", "flag", "uv"),
     dropout=("Dropout", "option", "d", float),
+    batch_size=("Number of words per training batch", "option", "bs", int),
+    max_length=("Max words per example.", "option", "xw", int),
+    min_length=("Min words per example.", "option", "nw", int),
     seed=("Seed for random number generators", "option", "s", float),
     nr_iter=("Number of iterations to pretrain", "option", "i", int),
 )
@@ -42,6 +45,9 @@
     use_vectors=False,
     dropout=0.2,
     nr_iter=1000,
+    batch_size=3000,
+    max_length=500,
+    min_length=5,
     seed=0,
 ):
     """
@@ -109,9 +115,14 @@
     msg.row(("#", "# Words", "Total Loss", "Loss", "w/s"), **row_settings)
     for epoch in range(nr_iter):
         for batch in util.minibatch_by_words(
-            ((text, None) for text in texts), size=3000
+            ((text, None) for text in texts), size=batch_size
         ):
-            docs = make_docs(nlp, [text for (text, _) in batch])
+            docs = make_docs(
+                nlp,
+                [text for (text, _) in batch],
+                max_length=max_length,
+                min_length=min_length,
+            )
             loss = make_update(model, docs, optimizer, drop=dropout)
             progress = tracker.update(epoch, loss, docs)
             if progress:
@@ -152,7 +163,7 @@
     return float(loss)
 
 
-def make_docs(nlp, batch, min_length=1, max_length=500):
+def make_docs(nlp, batch, min_length, max_length):
     docs = []
     for record in batch:
         text = record["text"]
@@ -241,11 +252,23 @@
             status = (
                 epoch,
                 self.nr_word,
-                "%.8f" % self.loss,
-                "%.8f" % loss_per_word,
+                _smart_round(self.loss, width=10),
+                _smart_round(loss_per_word, width=6),
                 int(wps),
             )
             self.prev_loss = float(self.loss)
             return status
         else:
             return None
+
+
+def _smart_round(figure, width=10, max_decimal=4):
+    """Round large numbers as integers, smaller numbers as decimals."""
+    n_digits = len(str(int(figure)))
+    n_decimal = width - (n_digits + 1)
+    if n_decimal <= 1:
+        return str(int(figure))
+    else:
+        n_decimal = min(n_decimal, max_decimal)
+        format_str = "%." + str(n_decimal) + "f"
+        return format_str % figure

{"golden_diff": "diff --git a/spacy/cli/pretrain.py b/spacy/cli/pretrain.py\n--- a/spacy/cli/pretrain.py\n+++ b/spacy/cli/pretrain.py\n@@ -29,6 +29,9 @@\n     embed_rows=(\"Embedding rows\", \"option\", \"er\", int),\n     use_vectors=(\"Whether to use the static vectors as input features\", \"flag\", \"uv\"),\n     dropout=(\"Dropout\", \"option\", \"d\", float),\n+    batch_size=(\"Number of words per training batch\", \"option\", \"bs\", int),\n+    max_length=(\"Max words per example.\", \"option\", \"xw\", int),\n+    min_length=(\"Min words per example.\", \"option\", \"nw\", int),\n     seed=(\"Seed for random number generators\", \"option\", \"s\", float),\n     nr_iter=(\"Number of iterations to pretrain\", \"option\", \"i\", int),\n )\n@@ -42,6 +45,9 @@\n     use_vectors=False,\n     dropout=0.2,\n     nr_iter=1000,\n+    batch_size=3000,\n+    max_length=500,\n+    min_length=5,\n     seed=0,\n ):\n     \"\"\"\n@@ -109,9 +115,14 @@\n     msg.row((\"#\", \"# Words\", \"Total Loss\", \"Loss\", \"w/s\"), **row_settings)\n     for epoch in range(nr_iter):\n         for batch in util.minibatch_by_words(\n-            ((text, None) for text in texts), size=3000\n+            ((text, None) for text in texts), size=batch_size\n         ):\n-            docs = make_docs(nlp, [text for (text, _) in batch])\n+            docs = make_docs(\n+                nlp,\n+                [text for (text, _) in batch],\n+                max_length=max_length,\n+                min_length=min_length,\n+            )\n             loss = make_update(model, docs, optimizer, drop=dropout)\n             progress = tracker.update(epoch, loss, docs)\n             if progress:\n@@ -152,7 +163,7 @@\n     return float(loss)\n \n \n-def make_docs(nlp, batch, min_length=1, max_length=500):\n+def make_docs(nlp, batch, min_length, max_length):\n     docs = []\n     for record in batch:\n         text = record[\"text\"]\n@@ -241,11 +252,23 @@\n             status = (\n                 epoch,\n                 self.nr_word,\n-                \"%.8f\" % self.loss,\n-                \"%.8f\" % loss_per_word,\n+                _smart_round(self.loss, width=10),\n+                _smart_round(loss_per_word, width=6),\n                 int(wps),\n             )\n             self.prev_loss = float(self.loss)\n             return status\n         else:\n             return None\n+\n+\n+def _smart_round(figure, width=10, max_decimal=4):\n+    \"\"\"Round large numbers as integers, smaller numbers as decimals.\"\"\"\n+    n_digits = len(str(int(figure)))\n+    n_decimal = width - (n_digits + 1)\n+    if n_decimal <= 1:\n+        return str(int(figure))\n+    else:\n+        n_decimal = min(n_decimal, max_decimal)\n+        format_str = \"%.\" + str(n_decimal) + \"f\"\n+        return format_str % figure\n", "issue": "Pretraining: batch size & document length\n## Feature description\r\n\r\nI was just testing the `pretrain` command and ran out of memory ^^ !\r\nLooking at the code, it seems that the `batch_size` is fixed to 3000. I also noticed that document with length greater than 500 token are pruned by default. I guess this are default parameters for a corpus of sentence / small text... But maybe we should let the user set these parameters ?\r\n\r\n\n", "before_files": [{"content": "# coding: utf8\nfrom __future__ import print_function, unicode_literals\n\nimport plac\nimport random\nimport numpy\nimport time\nfrom collections import Counter\nfrom pathlib import Path\nfrom thinc.v2v import Affine, Maxout\nfrom thinc.misc import LayerNorm as LN\nfrom thinc.neural.util import prefer_gpu\nfrom wasabi import Printer\nimport srsly\n\nfrom ..tokens import Doc\nfrom ..attrs import ID, HEAD\nfrom .._ml import Tok2Vec, flatten, chain, create_default_optimizer\nfrom .._ml import masked_language_model\nfrom .. import util\n\n\n@plac.annotations(\n    texts_loc=(\"Path to jsonl file with texts to learn from\", \"positional\", None, str),\n    vectors_model=(\"Name or path to vectors model to learn from\"),\n    output_dir=(\"Directory to write models each epoch\", \"positional\", None, str),\n    width=(\"Width of CNN layers\", \"option\", \"cw\", int),\n    depth=(\"Depth of CNN layers\", \"option\", \"cd\", int),\n    embed_rows=(\"Embedding rows\", \"option\", \"er\", int),\n    use_vectors=(\"Whether to use the static vectors as input features\", \"flag\", \"uv\"),\n    dropout=(\"Dropout\", \"option\", \"d\", float),\n    seed=(\"Seed for random number generators\", \"option\", \"s\", float),\n    nr_iter=(\"Number of iterations to pretrain\", \"option\", \"i\", int),\n)\ndef pretrain(\n    texts_loc,\n    vectors_model,\n    output_dir,\n    width=96,\n    depth=4,\n    embed_rows=2000,\n    use_vectors=False,\n    dropout=0.2,\n    nr_iter=1000,\n    seed=0,\n):\n    \"\"\"\n    Pre-train the 'token-to-vector' (tok2vec) layer of pipeline components,\n    using an approximate language-modelling objective. Specifically, we load\n    pre-trained vectors, and train a component like a CNN, BiLSTM, etc to predict\n    vectors which match the pre-trained ones. The weights are saved to a directory\n    after each epoch. You can then pass a path to one of these pre-trained weights\n    files to the 'spacy train' command.\n\n    This technique may be especially helpful if you have little labelled data.\n    However, it's still quite experimental, so your mileage may vary.\n\n    To load the weights back in during 'spacy train', you need to ensure\n    all settings are the same between pretraining and training. The API and\n    errors around this need some improvement.\n    \"\"\"\n    config = dict(locals())\n    msg = Printer()\n    util.fix_random_seed(seed)\n\n    has_gpu = prefer_gpu()\n    msg.info(\"Using GPU\" if has_gpu else \"Not using GPU\")\n\n    output_dir = Path(output_dir)\n    if not output_dir.exists():\n        output_dir.mkdir()\n        msg.good(\"Created output directory\")\n    srsly.write_json(output_dir / \"config.json\", config)\n    msg.good(\"Saved settings to config.json\")\n\n    # Load texts from file or stdin\n    if texts_loc != \"-\":  # reading from a file\n        texts_loc = Path(texts_loc)\n        if not texts_loc.exists():\n            msg.fail(\"Input text file doesn't exist\", texts_loc, exits=1)\n        with msg.loading(\"Loading input texts...\"):\n            texts = list(srsly.read_jsonl(texts_loc))\n        msg.good(\"Loaded input texts\")\n        random.shuffle(texts)\n    else:  # reading from stdin\n        msg.text(\"Reading input text from stdin...\")\n        texts = srsly.read_jsonl(\"-\")\n\n    with msg.loading(\"Loading model '{}'...\".format(vectors_model)):\n        nlp = util.load_model(vectors_model)\n    msg.good(\"Loaded model '{}'\".format(vectors_model))\n    pretrained_vectors = None if not use_vectors else nlp.vocab.vectors.name\n    model = create_pretraining_model(\n        nlp,\n        Tok2Vec(\n            width,\n            embed_rows,\n            conv_depth=depth,\n            pretrained_vectors=pretrained_vectors,\n            bilstm_depth=0,  # Requires PyTorch. Experimental.\n            cnn_maxout_pieces=3,  # You can try setting this higher\n            subword_features=True,  # Set to False for Chinese etc\n        ),\n    )\n    optimizer = create_default_optimizer(model.ops)\n    tracker = ProgressTracker(frequency=10000)\n    msg.divider(\"Pre-training tok2vec layer\")\n    row_settings = {\"widths\": (3, 10, 10, 6, 4), \"aligns\": (\"r\", \"r\", \"r\", \"r\", \"r\")}\n    msg.row((\"#\", \"# Words\", \"Total Loss\", \"Loss\", \"w/s\"), **row_settings)\n    for epoch in range(nr_iter):\n        for batch in util.minibatch_by_words(\n            ((text, None) for text in texts), size=3000\n        ):\n            docs = make_docs(nlp, [text for (text, _) in batch])\n            loss = make_update(model, docs, optimizer, drop=dropout)\n            progress = tracker.update(epoch, loss, docs)\n            if progress:\n                msg.row(progress, **row_settings)\n                if texts_loc == \"-\" and tracker.words_per_epoch[epoch] >= 10 ** 7:\n                    break\n        with model.use_params(optimizer.averages):\n            with (output_dir / (\"model%d.bin\" % epoch)).open(\"wb\") as file_:\n                file_.write(model.tok2vec.to_bytes())\n            log = {\n                \"nr_word\": tracker.nr_word,\n                \"loss\": tracker.loss,\n                \"epoch_loss\": tracker.epoch_loss,\n                \"epoch\": epoch,\n            }\n            with (output_dir / \"log.jsonl\").open(\"a\") as file_:\n                file_.write(srsly.json_dumps(log) + \"\\n\")\n        tracker.epoch_loss = 0.0\n        if texts_loc != \"-\":\n            # Reshuffle the texts if texts were loaded from a file\n            random.shuffle(texts)\n\n\ndef make_update(model, docs, optimizer, drop=0.0, objective=\"L2\"):\n    \"\"\"Perform an update over a single batch of documents.\n\n    docs (iterable): A batch of `Doc` objects.\n    drop (float): The droput rate.\n    optimizer (callable): An optimizer.\n    RETURNS loss: A float for the loss.\n    \"\"\"\n    predictions, backprop = model.begin_update(docs, drop=drop)\n    loss, gradients = get_vectors_loss(model.ops, docs, predictions, objective)\n    backprop(gradients, sgd=optimizer)\n    # Don't want to return a cupy object here\n    # The gradients are modified in-place by the BERT MLM,\n    # so we get an accurate loss\n    return float(loss)\n\n\ndef make_docs(nlp, batch, min_length=1, max_length=500):\n    docs = []\n    for record in batch:\n        text = record[\"text\"]\n        if \"tokens\" in record:\n            doc = Doc(nlp.vocab, words=record[\"tokens\"])\n        else:\n            doc = nlp.make_doc(text)\n        if \"heads\" in record:\n            heads = record[\"heads\"]\n            heads = numpy.asarray(heads, dtype=\"uint64\")\n            heads = heads.reshape((len(doc), 1))\n            doc = doc.from_array([HEAD], heads)\n        if len(doc) >= min_length and len(doc) < max_length:\n            docs.append(doc)\n    return docs\n\n\ndef get_vectors_loss(ops, docs, prediction, objective=\"L2\"):\n    \"\"\"Compute a mean-squared error loss between the documents' vectors and\n    the prediction.\n\n    Note that this is ripe for customization! We could compute the vectors\n    in some other word, e.g. with an LSTM language model, or use some other\n    type of objective.\n    \"\"\"\n    # The simplest way to implement this would be to vstack the\n    # token.vector values, but that's a bit inefficient, especially on GPU.\n    # Instead we fetch the index into the vectors table for each of our tokens,\n    # and look them up all at once. This prevents data copying.\n    ids = ops.flatten([doc.to_array(ID).ravel() for doc in docs])\n    target = docs[0].vocab.vectors.data[ids]\n    if objective == \"L2\":\n        d_scores = prediction - target\n        loss = (d_scores ** 2).sum()\n    else:\n        raise NotImplementedError(objective)\n    return loss, d_scores\n\n\ndef create_pretraining_model(nlp, tok2vec):\n    \"\"\"Define a network for the pretraining. We simply add an output layer onto\n    the tok2vec input model. The tok2vec input model needs to be a model that\n    takes a batch of Doc objects (as a list), and returns a list of arrays.\n    Each array in the output needs to have one row per token in the doc.\n    \"\"\"\n    output_size = nlp.vocab.vectors.data.shape[1]\n    output_layer = chain(\n        LN(Maxout(300, pieces=3)), Affine(output_size, drop_factor=0.0)\n    )\n    # This is annoying, but the parser etc have the flatten step after\n    # the tok2vec. To load the weights in cleanly, we need to match\n    # the shape of the models' components exactly. So what we cann\n    # \"tok2vec\" has to be the same set of processes as what the components do.\n    tok2vec = chain(tok2vec, flatten)\n    model = chain(tok2vec, output_layer)\n    model = masked_language_model(nlp.vocab, model)\n    model.tok2vec = tok2vec\n    model.output_layer = output_layer\n    model.begin_training([nlp.make_doc(\"Give it a doc to infer shapes\")])\n    return model\n\n\nclass ProgressTracker(object):\n    def __init__(self, frequency=1000000):\n        self.loss = 0.0\n        self.prev_loss = 0.0\n        self.nr_word = 0\n        self.words_per_epoch = Counter()\n        self.frequency = frequency\n        self.last_time = time.time()\n        self.last_update = 0\n        self.epoch_loss = 0.0\n\n    def update(self, epoch, loss, docs):\n        self.loss += loss\n        self.epoch_loss += loss\n        words_in_batch = sum(len(doc) for doc in docs)\n        self.words_per_epoch[epoch] += words_in_batch\n        self.nr_word += words_in_batch\n        words_since_update = self.nr_word - self.last_update\n        if words_since_update >= self.frequency:\n            wps = words_since_update / (time.time() - self.last_time)\n            self.last_update = self.nr_word\n            self.last_time = time.time()\n            loss_per_word = self.loss - self.prev_loss\n            status = (\n                epoch,\n                self.nr_word,\n                \"%.8f\" % self.loss,\n                \"%.8f\" % loss_per_word,\n                int(wps),\n            )\n            self.prev_loss = float(self.loss)\n            return status\n        else:\n            return None\n", "path": "spacy/cli/pretrain.py"}], "after_files": [{"content": "# coding: utf8\nfrom __future__ import print_function, unicode_literals\n\nimport plac\nimport random\nimport numpy\nimport time\nfrom collections import Counter\nfrom pathlib import Path\nfrom thinc.v2v import Affine, Maxout\nfrom thinc.misc import LayerNorm as LN\nfrom thinc.neural.util import prefer_gpu\nfrom wasabi import Printer\nimport srsly\n\nfrom ..tokens import Doc\nfrom ..attrs import ID, HEAD\nfrom .._ml import Tok2Vec, flatten, chain, create_default_optimizer\nfrom .._ml import masked_language_model\nfrom .. import util\n\n\n@plac.annotations(\n    texts_loc=(\"Path to jsonl file with texts to learn from\", \"positional\", None, str),\n    vectors_model=(\"Name or path to vectors model to learn from\"),\n    output_dir=(\"Directory to write models each epoch\", \"positional\", None, str),\n    width=(\"Width of CNN layers\", \"option\", \"cw\", int),\n    depth=(\"Depth of CNN layers\", \"option\", \"cd\", int),\n    embed_rows=(\"Embedding rows\", \"option\", \"er\", int),\n    use_vectors=(\"Whether to use the static vectors as input features\", \"flag\", \"uv\"),\n    dropout=(\"Dropout\", \"option\", \"d\", float),\n    batch_size=(\"Number of words per training batch\", \"option\", \"bs\", int),\n    max_length=(\"Max words per example.\", \"option\", \"xw\", int),\n    min_length=(\"Min words per example.\", \"option\", \"nw\", int),\n    seed=(\"Seed for random number generators\", \"option\", \"s\", float),\n    nr_iter=(\"Number of iterations to pretrain\", \"option\", \"i\", int),\n)\ndef pretrain(\n    texts_loc,\n    vectors_model,\n    output_dir,\n    width=96,\n    depth=4,\n    embed_rows=2000,\n    use_vectors=False,\n    dropout=0.2,\n    nr_iter=1000,\n    batch_size=3000,\n    max_length=500,\n    min_length=5,\n    seed=0,\n):\n    \"\"\"\n    Pre-train the 'token-to-vector' (tok2vec) layer of pipeline components,\n    using an approximate language-modelling objective. Specifically, we load\n    pre-trained vectors, and train a component like a CNN, BiLSTM, etc to predict\n    vectors which match the pre-trained ones. The weights are saved to a directory\n    after each epoch. You can then pass a path to one of these pre-trained weights\n    files to the 'spacy train' command.\n\n    This technique may be especially helpful if you have little labelled data.\n    However, it's still quite experimental, so your mileage may vary.\n\n    To load the weights back in during 'spacy train', you need to ensure\n    all settings are the same between pretraining and training. The API and\n    errors around this need some improvement.\n    \"\"\"\n    config = dict(locals())\n    msg = Printer()\n    util.fix_random_seed(seed)\n\n    has_gpu = prefer_gpu()\n    msg.info(\"Using GPU\" if has_gpu else \"Not using GPU\")\n\n    output_dir = Path(output_dir)\n    if not output_dir.exists():\n        output_dir.mkdir()\n        msg.good(\"Created output directory\")\n    srsly.write_json(output_dir / \"config.json\", config)\n    msg.good(\"Saved settings to config.json\")\n\n    # Load texts from file or stdin\n    if texts_loc != \"-\":  # reading from a file\n        texts_loc = Path(texts_loc)\n        if not texts_loc.exists():\n            msg.fail(\"Input text file doesn't exist\", texts_loc, exits=1)\n        with msg.loading(\"Loading input texts...\"):\n            texts = list(srsly.read_jsonl(texts_loc))\n        msg.good(\"Loaded input texts\")\n        random.shuffle(texts)\n    else:  # reading from stdin\n        msg.text(\"Reading input text from stdin...\")\n        texts = srsly.read_jsonl(\"-\")\n\n    with msg.loading(\"Loading model '{}'...\".format(vectors_model)):\n        nlp = util.load_model(vectors_model)\n    msg.good(\"Loaded model '{}'\".format(vectors_model))\n    pretrained_vectors = None if not use_vectors else nlp.vocab.vectors.name\n    model = create_pretraining_model(\n        nlp,\n        Tok2Vec(\n            width,\n            embed_rows,\n            conv_depth=depth,\n            pretrained_vectors=pretrained_vectors,\n            bilstm_depth=0,  # Requires PyTorch. Experimental.\n            cnn_maxout_pieces=3,  # You can try setting this higher\n            subword_features=True,  # Set to False for Chinese etc\n        ),\n    )\n    optimizer = create_default_optimizer(model.ops)\n    tracker = ProgressTracker(frequency=10000)\n    msg.divider(\"Pre-training tok2vec layer\")\n    row_settings = {\"widths\": (3, 10, 10, 6, 4), \"aligns\": (\"r\", \"r\", \"r\", \"r\", \"r\")}\n    msg.row((\"#\", \"# Words\", \"Total Loss\", \"Loss\", \"w/s\"), **row_settings)\n    for epoch in range(nr_iter):\n        for batch in util.minibatch_by_words(\n            ((text, None) for text in texts), size=batch_size\n        ):\n            docs = make_docs(\n                nlp,\n                [text for (text, _) in batch],\n                max_length=max_length,\n                min_length=min_length,\n            )\n            loss = make_update(model, docs, optimizer, drop=dropout)\n            progress = tracker.update(epoch, loss, docs)\n            if progress:\n                msg.row(progress, **row_settings)\n                if texts_loc == \"-\" and tracker.words_per_epoch[epoch] >= 10 ** 7:\n                    break\n        with model.use_params(optimizer.averages):\n            with (output_dir / (\"model%d.bin\" % epoch)).open(\"wb\") as file_:\n                file_.write(model.tok2vec.to_bytes())\n            log = {\n                \"nr_word\": tracker.nr_word,\n                \"loss\": tracker.loss,\n                \"epoch_loss\": tracker.epoch_loss,\n                \"epoch\": epoch,\n            }\n            with (output_dir / \"log.jsonl\").open(\"a\") as file_:\n                file_.write(srsly.json_dumps(log) + \"\\n\")\n        tracker.epoch_loss = 0.0\n        if texts_loc != \"-\":\n            # Reshuffle the texts if texts were loaded from a file\n            random.shuffle(texts)\n\n\ndef make_update(model, docs, optimizer, drop=0.0, objective=\"L2\"):\n    \"\"\"Perform an update over a single batch of documents.\n\n    docs (iterable): A batch of `Doc` objects.\n    drop (float): The droput rate.\n    optimizer (callable): An optimizer.\n    RETURNS loss: A float for the loss.\n    \"\"\"\n    predictions, backprop = model.begin_update(docs, drop=drop)\n    loss, gradients = get_vectors_loss(model.ops, docs, predictions, objective)\n    backprop(gradients, sgd=optimizer)\n    # Don't want to return a cupy object here\n    # The gradients are modified in-place by the BERT MLM,\n    # so we get an accurate loss\n    return float(loss)\n\n\ndef make_docs(nlp, batch, min_length, max_length):\n    docs = []\n    for record in batch:\n        text = record[\"text\"]\n        if \"tokens\" in record:\n            doc = Doc(nlp.vocab, words=record[\"tokens\"])\n        else:\n            doc = nlp.make_doc(text)\n        if \"heads\" in record:\n            heads = record[\"heads\"]\n            heads = numpy.asarray(heads, dtype=\"uint64\")\n            heads = heads.reshape((len(doc), 1))\n            doc = doc.from_array([HEAD], heads)\n        if len(doc) >= min_length and len(doc) < max_length:\n            docs.append(doc)\n    return docs\n\n\ndef get_vectors_loss(ops, docs, prediction, objective=\"L2\"):\n    \"\"\"Compute a mean-squared error loss between the documents' vectors and\n    the prediction.\n\n    Note that this is ripe for customization! We could compute the vectors\n    in some other word, e.g. with an LSTM language model, or use some other\n    type of objective.\n    \"\"\"\n    # The simplest way to implement this would be to vstack the\n    # token.vector values, but that's a bit inefficient, especially on GPU.\n    # Instead we fetch the index into the vectors table for each of our tokens,\n    # and look them up all at once. This prevents data copying.\n    ids = ops.flatten([doc.to_array(ID).ravel() for doc in docs])\n    target = docs[0].vocab.vectors.data[ids]\n    if objective == \"L2\":\n        d_scores = prediction - target\n        loss = (d_scores ** 2).sum()\n    else:\n        raise NotImplementedError(objective)\n    return loss, d_scores\n\n\ndef create_pretraining_model(nlp, tok2vec):\n    \"\"\"Define a network for the pretraining. We simply add an output layer onto\n    the tok2vec input model. The tok2vec input model needs to be a model that\n    takes a batch of Doc objects (as a list), and returns a list of arrays.\n    Each array in the output needs to have one row per token in the doc.\n    \"\"\"\n    output_size = nlp.vocab.vectors.data.shape[1]\n    output_layer = chain(\n        LN(Maxout(300, pieces=3)), Affine(output_size, drop_factor=0.0)\n    )\n    # This is annoying, but the parser etc have the flatten step after\n    # the tok2vec. To load the weights in cleanly, we need to match\n    # the shape of the models' components exactly. So what we cann\n    # \"tok2vec\" has to be the same set of processes as what the components do.\n    tok2vec = chain(tok2vec, flatten)\n    model = chain(tok2vec, output_layer)\n    model = masked_language_model(nlp.vocab, model)\n    model.tok2vec = tok2vec\n    model.output_layer = output_layer\n    model.begin_training([nlp.make_doc(\"Give it a doc to infer shapes\")])\n    return model\n\n\nclass ProgressTracker(object):\n    def __init__(self, frequency=1000000):\n        self.loss = 0.0\n        self.prev_loss = 0.0\n        self.nr_word = 0\n        self.words_per_epoch = Counter()\n        self.frequency = frequency\n        self.last_time = time.time()\n        self.last_update = 0\n        self.epoch_loss = 0.0\n\n    def update(self, epoch, loss, docs):\n        self.loss += loss\n        self.epoch_loss += loss\n        words_in_batch = sum(len(doc) for doc in docs)\n        self.words_per_epoch[epoch] += words_in_batch\n        self.nr_word += words_in_batch\n        words_since_update = self.nr_word - self.last_update\n        if words_since_update >= self.frequency:\n            wps = words_since_update / (time.time() - self.last_time)\n            self.last_update = self.nr_word\n            self.last_time = time.time()\n            loss_per_word = self.loss - self.prev_loss\n            status = (\n                epoch,\n                self.nr_word,\n                _smart_round(self.loss, width=10),\n                _smart_round(loss_per_word, width=6),\n                int(wps),\n            )\n            self.prev_loss = float(self.loss)\n            return status\n        else:\n            return None\n\n\ndef _smart_round(figure, width=10, max_decimal=4):\n    \"\"\"Round large numbers as integers, smaller numbers as decimals.\"\"\"\n    n_digits = len(str(int(figure)))\n    n_decimal = width - (n_digits + 1)\n    if n_decimal <= 1:\n        return str(int(figure))\n    else:\n        n_decimal = min(n_decimal, max_decimal)\n        format_str = \"%.\" + str(n_decimal) + \"f\"\n        return format_str % figure\n", "path": "spacy/cli/pretrain.py"}]}