!pip install "thinc>=8.0.0" transformers torch "ml_datasets>=0.2.0" "tqdm>=4.41"

"""
Explanation: Training a part-of-speech tagger with transformers (BERT)
This example shows how to use Thinc and Hugging Face's transformers library to implement and train a part-of-speech tagger on the Universal Dependencies AnCora corpus. This notebook assumes familiarity with machine learning concepts, transformer models and Thinc's config system and Model API (see the "Thinc for beginners" notebook and the documentation for more info).
End of explanation
"""


from thinc.api import prefer_gpu, use_pytorch_for_gpu_memory

is_gpu = prefer_gpu()
print("GPU:", is_gpu)
if is_gpu:
    use_pytorch_for_gpu_memory()

"""
Explanation: First, let's use Thinc's prefer_gpu helper to make sure we're performing operations on GPU if available. The function should be called right after importing Thinc, and it returns a boolean indicating whether the GPU has been activated. If we're on GPU, we can also call use_pytorch_for_gpu_memory to route cupy's memory allocation via PyTorch, so both can play together nicely.
End of explanation
"""


CONFIG = """
[model]
@layers = "TransformersTagger.v1"
starter = "bert-base-multilingual-cased"

[optimizer]
@optimizers = "Adam.v1"

[optimizer.learn_rate]
@schedules = "warmup_linear.v1"
initial_rate = 0.01
warmup_steps = 3000
total_steps = 6000

[loss]
@losses = "SequenceCategoricalCrossentropy.v1"

[training]
batch_size = 128
words_per_subbatch = 2000
n_epoch = 10
"""

"""
Explanation: Overview: the final config
Here's the final config for the model we're building in this notebook. It references a custom TransformersTagger that takes the name of a starter (the pretrained model to use), an optimizer, a learning rate schedule with warm-up and the general training settings. You can keep the config string within your file or notebook, or save it to a conig.cfg file and load it in via Config.from_disk.
End of explanation
"""


from typing import Optional, List
import numpy
from thinc.types import Ints1d, Floats2d
from dataclasses import dataclass
import torch
from transformers import BatchEncoding, TokenSpan


@dataclass
class TokensPlus:
    batch_size: int
    tok2wp: List[Ints1d]
    input_ids: torch.Tensor
    token_type_ids: torch.Tensor
    attention_mask: torch.Tensor
        
    def __init__(self, inputs: List[List[str]], wordpieces: BatchEncoding):
        self.input_ids = wordpieces["input_ids"]
        self.attention_mask = wordpieces["attention_mask"]
        self.token_type_ids = wordpieces["token_type_ids"]
        self.batch_size = self.input_ids.shape[0]
        self.tok2wp = []
        for i in range(self.batch_size):
            spans = [wordpieces.word_to_tokens(i, j) for j in range(len(inputs[i]))]
            self.tok2wp.append(self.get_wp_starts(spans))
        
    def get_wp_starts(self, spans: List[Optional[TokenSpan]]) -> Ints1d:
        """Calculate an alignment mapping each token index to its first wordpiece."""
        alignment = numpy.zeros((len(spans)), dtype="i")
        for i, span in enumerate(spans):
            if span is None:
                raise ValueError(
                    "Token did not align to any wordpieces. Was the tokenizer "
                    "run with is_split_into_words=True?"
                )
            else:
                alignment[i] = span.start
        return alignment
    

def test_tokens_plus(name: str="bert-base-multilingual-cased"):
    from transformers import AutoTokenizer
    inputs = [
        ["Our", "band", "is", "called", "worlthatmustbedivided", "!"],
        ["We", "rock", "!"]
    ]
    tokenizer = AutoTokenizer.from_pretrained(name)
    wordpieces = tokenizer(
        inputs,
        is_split_into_words=True,
        add_special_tokens=True,
        return_token_type_ids=True,
        return_attention_mask=True,
        return_length=True,
        return_tensors="pt",
        padding="longest"
    )
    tplus = TokensPlus(inputs, wordpieces)
    assert len(tplus.tok2wp) == len(inputs) == len(tplus.input_ids)
    for i, align in enumerate(tplus.tok2wp):
        assert len(align) == len(inputs[i])
        for j in align:
            assert j >= 0 and j < tplus.input_ids.shape[1]
            
test_tokens_plus()

"""
Explanation: Defining the model
The Thinc model we want to define should consist of 3 components: the transformers tokenizer, the actual transformer implemented in PyTorch and a softmax-activated output layer.
1. Wrapping the tokenizer
To make it easier to keep track of the data that's passed around (and get type errors if something goes wrong), we first create a TokensPlus dataclass that holds the information we need from the transformers tokenizer. The most important work we'll do in this class is to build an alignment map. The transformer models are trained on input sequences that over-segment the sentence, so that they can work on smaller vocabularies. These over-segmentations are generally called "word pieces". The transformer will return a tensor with one vector per wordpiece. We need to map that to a tensor with one vector per POS-tagged token. We'll pass those token representations into a feed-forward network to predict the tag probabilities. During the backward pass, we'll then need to invert this mapping, so that we can calculate the gradients with respect to the wordpieces given the gradients with respect to the tokens. To keep things relatively simple, we'll store the alignment as a list of arrays, with each array mapping one token to one wordpiece vector (its first one). To make this work, we'll need to run the tokenizer with is_split_into_words=True, which should ensure that we get at least one wordpiece per token.
End of explanation
"""


import thinc
from thinc.api import Model
from transformers import AutoTokenizer

@thinc.registry.layers("transformers_tokenizer.v1")
def TransformersTokenizer(name: str) -> Model[List[List[str]], TokensPlus]:
    def forward(model, inputs: List[List[str]], is_train: bool):
        tokenizer = model.attrs["tokenizer"]
        wordpieces = tokenizer(
            inputs,
            is_split_into_words=True,
            add_special_tokens=True,
            return_token_type_ids=True,
            return_attention_mask=True,
            return_length=True,
            return_tensors="pt",
            padding="longest"
        )
        return TokensPlus(inputs, wordpieces), lambda d_tokens: []

    return Model("tokenizer", forward, attrs={"tokenizer": AutoTokenizer.from_pretrained(name)})

"""
Explanation: The wrapped tokenizer will take a list-of-lists as input (the texts) and will output a TokensPlus object containing the fully padded batch of tokens. The wrapped transformer will take a list of TokensPlus objects and will output a list of 2-dimensional arrays.

TransformersTokenizer: List[List[str]] → TokensPlus
Transformer: TokensPlus → List[Array2d]


💡 Since we're adding type hints everywhere (and Thinc is fully typed, too), you can run your code through mypy to find type errors and inconsistencies. If you're using an editor like Visual Studio Code, you can enable mypy linting and type errors will be highlighted in real time as you write code.

To use the tokenizer as a layer in our network, we register a new function that returns a Thinc Model. The function takes the name of the pretrained weights (e.g. "bert-base-multilingual-cased") as an argument that can later be provided via the config. After loading the AutoTokenizer, we can stash it in the attributes. This lets us access it at any point later on via model.attrs["tokenizer"].
End of explanation
"""


from typing import List, Tuple, Callable
from thinc.api import ArgsKwargs, torch2xp, xp2torch
from thinc.types import Floats2d

def convert_transformer_inputs(model, tokens: TokensPlus, is_train):
    kwargs = {
        "input_ids": tokens.input_ids,
        "attention_mask": tokens.attention_mask,
        "token_type_ids": tokens.token_type_ids,
    }
    return ArgsKwargs(args=(), kwargs=kwargs), lambda dX: []


def convert_transformer_outputs(
    model: Model,
    inputs_outputs: Tuple[TokensPlus, Tuple[torch.Tensor]],
    is_train: bool
) -> Tuple[List[Floats2d], Callable]:
    tplus, trf_outputs = inputs_outputs
    wp_vectors = torch2xp(trf_outputs[0])
    tokvecs = [wp_vectors[i, idx] for i, idx in enumerate(tplus.tok2wp)]

    def backprop(d_tokvecs: List[Floats2d]) -> ArgsKwargs:
        # Restore entries for BOS and EOS markers
        d_wp_vectors = model.ops.alloc3f(*trf_outputs[0].shape, dtype="f")
        for i, idx in enumerate(tplus.tok2wp):
            d_wp_vectors[i, idx] += d_tokvecs[i]
        return ArgsKwargs(
            args=(trf_outputs[0],),
            kwargs={"grad_tensors": xp2torch(d_wp_vectors)},
        )

    return tokvecs, backprop

"""
Explanation: The forward pass takes the model and a list-of-lists of strings and outputs the TokensPlus dataclass. It also outputs a dummy callback function, to meet the API contract for Thinc models. Even though there's no way we can meaningfully "backpropagate" this layer, we need to make sure the function has the right signature, so that it can be used interchangeably with other layers.
2. Wrapping the transformer
To load and wrap the transformer, we can use transformers.AutoModel and Thinc's PyTorchWrapper. The forward method of the wrapped model can take arbitrary positional arguments and keyword arguments. Here's what the wrapped model is going to look like:
python
@thinc.registry.layers("transformers_model.v1")
def Transformer(name) -&gt; Model[TokensPlus, List[Floats2d]]:
    return PyTorchWrapper(
        AutoModel.from_pretrained(name),
        convert_inputs=convert_transformer_inputs,
        convert_outputs=convert_transformer_outputs,
    )
The Transformer layer takes our TokensPlus dataclass as input and outputs a list of 2-dimensional arrays. The convert functions are used to map inputs and outputs to and from the PyTorch model. Each function should return the converted output, and a callback to use during the backward pass. To make the arbitrary positional and keyword arguments easier to manage, Thinc uses an ArgsKwargs dataclass, essentially a named tuple with args and kwargs that can be spread into a function as *ArgsKwargs.args and **ArgsKwargs.kwargs. The ArgsKwargs objects will be passed straight into the model in the forward pass, and straight into torch.autograd.backward during the backward pass.
End of explanation
"""


import thinc
from thinc.api import PyTorchWrapper
from transformers import AutoModel

@thinc.registry.layers("transformers_model.v1")
def Transformer(name: str) -> Model[TokensPlus, List[Floats2d]]:
    return PyTorchWrapper(
        AutoModel.from_pretrained(name),
        convert_inputs=convert_transformer_inputs,
        convert_outputs=convert_transformer_outputs,
    )

"""
Explanation: The input and output transformation functions give you full control of how data is passed into and out of the underlying PyTorch model, so you can work with PyTorch layers that expect and return arbitrary objects. Putting it all together, we now have a nice layer that is configured with the name of a transformer model, that acts as a function mapping tokenized input into feature vectors.
End of explanation
"""


from thinc.api import chain, with_array, Softmax

@thinc.registry.layers("TransformersTagger.v1")
def TransformersTagger(starter: str, n_tags: int = 17) -> Model[List[List[str]], List[Floats2d]]:
    return chain(
        TransformersTokenizer(starter),
        Transformer(starter),
        with_array(Softmax(n_tags)),
    )

"""
Explanation: We can now combine the TransformersTokenizer and Transformer into a feed-forward network using the chain combinator. The with_array layer transforms a sequence of data into a contiguous 2d array on the way into and
out of a model.
End of explanation
"""


from thinc.api import Config, registry

C = registry.resolve(Config().from_str(CONFIG))

model = C["model"]
optimizer = C["optimizer"]
calculate_loss = C["loss"]
cfg = C["training"]

"""
Explanation: Training the model
Setting up model and data
Since we've registered all layers via @thinc.registry.layers, we can construct the model, its settings and other functions we need from a config (see CONFIG above). The result is a config object with a model, an optimizer, a function to calculate the loss and the training settings.
End of explanation
"""


import ml_datasets
(train_X, train_Y), (dev_X, dev_Y) = ml_datasets.ud_ancora_pos_tags()

train_Y = list(map(model.ops.asarray, train_Y))  # convert to cupy if needed
dev_Y = list(map(model.ops.asarray, dev_Y))  # convert to cupy if needed

model.initialize(X=train_X[:5], Y=train_Y[:5])

"""
Explanation: We’ve prepared a separate package ml-datasets with loaders for some common datasets, including the AnCora data. If we're using a GPU, calling ops.asarray on the outputs ensures that they're converted to cupy arrays (instead of numpy arrays). Calling Model.initialize with a batch of inputs and outputs allows Thinc to infer the missing dimensions.
End of explanation
"""


def minibatch_by_words(pairs, max_words):
    pairs = list(zip(*pairs))
    pairs.sort(key=lambda xy: len(xy[0]), reverse=True)
    batch = []
    for X, Y in pairs:
        batch.append((X, Y))
        n_words = max(len(xy[0]) for xy in batch) * len(batch)
        if n_words >= max_words:
            yield batch[:-1]
            batch = [(X, Y)]
    if batch:
        yield batch

def evaluate_sequences(model, Xs: List[Floats2d], Ys: List[Floats2d], batch_size: int) -> float:
    correct = 0.0
    total = 0.0
    for X, Y in model.ops.multibatch(batch_size, Xs, Ys):
        Yh = model.predict(X)
        for yh, y in zip(Yh, Y):
            correct += (y.argmax(axis=1) == yh.argmax(axis=1)).sum()
            total += y.shape[0]
    return float(correct / total)

"""
Explanation: Helper functions for training and evaluation
Before we can train the model, we also need to set up the following helper functions for batching and evaluation:

minibatch_by_words: Group pairs of sequences into minibatches under max_words in size, considering padding. The size of a padded batch is the length of its longest sequence multiplied by the number of elements in the batch.
evaluate_sequences: Evaluate the model sequences of two-dimensional arrays and return the score.
End of explanation
"""


from tqdm.notebook import tqdm
from thinc.api import fix_random_seed

fix_random_seed(0)

for epoch in range(cfg["n_epoch"]):
    batches = model.ops.multibatch(cfg["batch_size"], train_X, train_Y, shuffle=True)
    for outer_batch in tqdm(batches, leave=False):
        for batch in minibatch_by_words(outer_batch, cfg["words_per_subbatch"]):
            inputs, truths = zip(*batch)
            inputs = list(inputs)
            guesses, backprop = model(inputs, is_train=True)
            backprop(calculate_loss.get_grad(guesses, truths))
        model.finish_update(optimizer)
        optimizer.step_schedules()
    score = evaluate_sequences(model, dev_X, dev_Y, cfg["batch_size"])
    print(epoch, f"{score:.3f}")

"""
Explanation: The training loop
Transformers often learn best with large batch sizes – larger than fits in GPU memory. But you don't have to backprop the whole batch at once. Here we consider the "logical" batch size (number of examples per update) separately from the physical batch size. For the physical batch size, what we care about is the number of words (considering padding too). We also want to sort by length, for efficiency. 
At the end of the batch, we call the optimizer with the accumulated gradients, and advance the learning rate schedules. You might want to evaluate more often than once per epoch – that's up to you.
End of explanation
"""


