!pip install "thinc>=8.0.0a0"

"""
Explanation: Benchmarking Thinc layers with a custom benchmark layer
This notebook shows how to write a benchmark layer that can wrap any layer(s) in your network and that logs the execution times of the initialization, forward pass and backward pass. The benchmark layer can also be mapped to an operator like @ to make it easy to add debugging to your network.
End of explanation
"""


import logging

logger = logging.getLogger("thinc:benchmark")
if not logger.hasHandlers():  # prevent Jupyter from adding multiple loggers
    formatter = logging.Formatter('%(asctime)s %(name)s  %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

"""
Explanation: To log the results, we first set up a custom logger using Python's logging module. You could also just print the stats instead, but using logging is cleaner, since it lets other users modify the logger's behavior more easily, and separates the logs from other output and write it to a file (e.g. if you're benchmarking several layers during training). The following logging config will output the date and time, the name of the logger and the logged results.
End of explanation
"""


from timeit import default_timer
from wasabi import color


class TimeLogger:
    def __init__(self, name):
        self.colors = {"forward": "green", "backward": "blue"}
        self.name = name
        self.timers = {}
        
    def start(self, name):
        self.timers[name] = default_timer()
        
    def end(self, name):
        result = default_timer() - self.timers[name]
        label = f"{name.upper():<8}"
        label = color(label, self.colors.get(name), bold=True)
        logger.debug(f"{self.name:<12} | {label} | {result:.6f}")

"""
Explanation: Here's a minimalistic time logger that can be initialized with the name of a given layer, and can track several events (e.g. "forward" and "backward"). When the TimeLogger.end method is called, the output is formatted nicely and the elapsed time is logged with the logger name and colored label.
End of explanation
"""


from thinc.api import Model
        
def benchmark(layer, name=None):
    name = name if name is not None else layer.name
    t = TimeLogger(name)
    
    def init(model, X, Y):
        t.start("init")
        result = layer.initialize(X, Y)
        t.end("init")
        return result
        
    def forward(model, X, is_train):
        t.start("forward")
        layer_Y, layer_callback = layer(X, is_train=is_train)
        t.end("forward")
        
        def backprop(dY):
            t.start("backward")
            result = layer_callback(dY)
            t.end("backward")
            return result
        
        return layer_Y, backprop
        
    return Model(f"benchmark:{layer.name}", forward, init=init) 

"""
Explanation: The benchmark layer now has to wrap the forward pass, backward pass and initialization of the layer it wraps and log the execution times. It then returns a Thinc model instance with the custom forward function and a custom init function. We'll also allow setting a custom name to make it easier to tell multiple wrapped benchmark layers apart.
End of explanation
"""


import numpy
from thinc.api import chain, Linear

X = numpy.zeros((1, 2), dtype="f")

model = benchmark(chain(benchmark(Linear(1)), Linear(1)), name="outer")
model.initialize(X=X)
Y, backprop = model(X, is_train=False)
dX = backprop(Y)

"""
Explanation: Usage examples
Using the benchmark layer as a function
We can now wrap one or more layers (including nested layers) with the benchmark function. This is the original model:
python
model = chain(Linear(1), Linear(1))
End of explanation
"""


from thinc.api import Model

with Model.define_operators({">>": chain, "@": benchmark}):
    model = (Linear(1) @ "first" >> Linear(1)) @ "outer"
    
model.initialize(X=X)
Y, backprop = model(X, is_train=True)
dX = backprop(Y)

"""
Explanation: Using the benchmark layer as an operator
Alternatively, we can also use Model.define_operators to map benchmark to an operator like @. The left argument of the operator is the first argument passed into the function (the layer) and the right argument is the second argument (the name). The following example wraps the whole network (two chained Linear layers) in a benchmark layer named "outer", and the first Linear layer in a benchmark layer named "first".
End of explanation
"""


from thinc.api import Model, chain, Relu, Softmax, Adam

n_hidden = 32
dropout = 0.2

with Model.define_operators({">>": chain, "@": benchmark}):
    model = (
        Relu(nO=n_hidden, dropout=dropout) @ "relu1"
        >> Relu(nO=n_hidden, dropout=dropout) @ "relu2"
        >> Softmax()
    )

train_X = numpy.zeros((5, 784), dtype="f")
train_Y = numpy.zeros((540, 10), dtype="f")

model.initialize(X=train_X[:5], Y=train_Y[:5])
optimizer = Adam(0.001)
for i in range(10):
    for X, Y in model.ops.multibatch(8, train_X, train_Y, shuffle=True):
        Yh, backprop = model.begin_update(X)
        backprop(Yh - Y)
        model.finish_update(optimizer)

"""
Explanation: Using the benchmark layer during training
End of explanation
"""


