# coding: utf8
from __future__ import unicode_literals, division, print_function

import plac
from pathlib import Path
import ujson
import cProfile
import pstats

import spacy
import sys
import tqdm
import cytoolz


def read_inputs(loc):
    if loc is None:
        file_ = sys.stdin
        file_ = (line.encode('utf8') for line in file_)
    else:
        file_ = Path(loc).open()
    for line in file_:
        data = ujson.loads(line)
        text = data['text']
        yield text


@plac.annotations(
    lang=("model/language", "positional", None, str),
    inputs=("Location of input file", "positional", None, read_inputs))
def profile(cmd, lang, inputs=None):
    """
    Profile a spaCy pipeline, to find out which functions take the most time.
    """
    nlp = spacy.load(lang)
    texts = list(cytoolz.take(10000, inputs))
    cProfile.runctx("parse_texts(nlp, texts)", globals(), locals(),
                    "Profile.prof")
    s = pstats.Stats("Profile.prof")
    s.strip_dirs().sort_stats("time").print_stats()


def parse_texts(nlp, texts):
    for doc in nlp.pipe(tqdm.tqdm(texts), batch_size=128):
        pass

@@ -207,16 +207,44 @@
 cytoolz%0A
+import thinc.extra.datasets%0A
 %0A%0Adef re
@@ -781,24 +781,167 @@
 me.%0A    %22%22%22%0A
+    if inputs is None:%0A        imdb_train, _ = thinc.extra.datasets.imdb()%0A        inputs, _ = zip(*imdb_train)%0A        inputs = inputs%5B:2000%5D%0A
     nlp = sp
@@ -1173,16 +1173,19 @@
 _stats(%22
+cum
 time%22).p
@@ -1286,10 +1286,9 @@
 ze=1
-28
+6
 ):%0A 

Improve profiling