Wrapping independently trained Pytorch model with Thinc

[kanayer]

I trained the POS tagging model for Korean in Pytorch separately from spaCy. Is it possible to wrap the trained model with PytorchWrapper and add it to the custom pipeline?

For now i did something like this:
from thinc.api import PyTorchWrapper
tagger = "./models/model_pt" tagger_model = PyTorchWrapper(tagger)

Can you please advise me on what to do next from here? How to add this model to the pipeline? Or should I re-train my model from scratch with spaCy?

[adrianeboyd]

To use this in spacy, you would need to write a custom pipeline component that can load this model and set the right annotation on the doc from the model predictions.

There's a similar example in a demo project here (it's for NER instead of POS, but the basics for the pipeline component should be similar):

https://github.com/explosion/projects/tree/v3/tutorials/ner_pytorch_medical

Look at the custom functions under scripts/. This example allows you to also train the model with spacy train, but I think you can write a simpler pipeline component that just loads the pytorch model from a path in initialize and implements predict and set_annotations for that model without implementing the training side of things.

[kanayer]

Thank you very much for your reply. I have made the following steps however I keep getting the error message after trying to use nlp.add_pipe. Could you please help me understand what am I doing wrong?

STEP 1: Registered custom new factory similar to the one in the tutorial:
@Language.factory("torch_tagger") def make_torch_tagger(nlp, name, model): return TorchPOSTagger(nlp.vocab, model, name)

STEP 2: Created custom pipeline component similar to the one in the tutorial, but including only initialize, predict, and set_annotations:

`class TorchPOSTagger(TrainablePipe):
def init(self, vocab, model, name: str = "torch_tagger"):

    self.model = model
    self.vocab = vocab
    self.name = name
    cfg = {"labels": []}
    self.cfg = dict(sorted(cfg.items()))

def labels(self):
    labels = [POS_TAG.vocab.stoi]
    return tuple(self.cfg["labels"])

def add_label(self, label):
    self.cfg["labels"] = list(self.labels) + [label]

def initialize(self, get_examples: Callable[[], Iterable[Example]], *, nlp=None, labels=None):
    
    validate_get_examples(get_examples, "TorchPOSTagger.initialize")
    if labels is not None:
        for tag in labels:
            self.add_label(tag)
    else:
        tags = set()
        for example in get_examples():
            for token in example.y:
                if token.ent_type_:
                    tags.add(token.ent_type_)
        for tag in sorted(tags):
            self.add_label(tag)
    doc_sample = []
    for example in islice(get_examples(), 10):
        doc_sample.append(example.x)

    self._require_labels()
    assert len(doc_sample) > 0, Errors.E923.format(name=self.name)
    self.model.initialize(X=doc_sample, Y=self.labels)
    nlp.config["components"][self.name]["model"]["nO"] = len(self.labels)

def predict(self, docs: Iterable[Doc]) -> Iterable[Ints1d]:
    if not any(len(doc) for doc in docs):
        # Handle cases where there are no tokens in any docs.
        n_labels = len(self.labels)
        guesses = [self.model.ops.alloc((0, n_labels)) for doc in docs]
        assert len(guesses) == len(docs)
        return guesses
    scores = self.model.predict(docs)

    assert len(scores) == len(docs), (len(scores), len(docs))
    guesses = []
    for doc_scores in scores:
        doc_guesses = doc_scores.argmax(axis=1)
        if not isinstance(doc_guesses, numpy.ndarray):
            doc_guesses = doc_guesses.get()
        guesses.append(doc_guesses)
    assert len(guesses) == len(docs)
    return guesses

def set_annotations(self, docs: Iterable[Doc], preds: Iterable[Ints1d]):
    
    if isinstance(docs, Doc):
        docs = [docs]
    for doc, tag_ids in zip(docs, preds):
        labels = iob_to_biluo([self.labels[tag_id] for tag_id in tag_ids])
        try:
            spans = biluo_tags_to_spans(doc, labels)
        except ValueError:
            spans = []
        doc.ents = spans`

STEP 3: Created nlp object for Korean
nlp = spacy.blank("ko")

STEP 4: Tried to add the model into pipeline by using nlp.add_pipe:
nlp.add_pipe("torch_tagger", tagger_model_wrapped).

However, I got the following error:

ConfigValidationError:

Config validation error

torch_tagger -> model field required

{'nlp':<spacy.lang.ko.Korean object at 0x7f05cdea3390>, 'name': <thinc.model.Model object at 0x7f05ce950ef0>, '@factories': 'torch_tagger'}

Even if I send the Thinc wrapped model as a parameter, the error message remains. Could you please navigate me in the right direction? Do I pass the model as a parameter? Also, do I need to create a config file?

[adrianeboyd]

This is a place where there's no a perfect fit in the current API. It's useful to inherit the __call__, pipe, and serialization methods from TrainablePipe, but you don't actually want to train your model (unless I'm misunderstanding), so you don't need this kind of initialization, which is meant for a blank thinc Model.

I think what you want would look more like this (this is untested typed off the top of my head!):

class TorchPOSTagger(TrainablePipe):
    def initialize(self, path):
        self.model = PyTorchWrapper(path, convert_inputs=, convert_outputs=) # however this works for your model

    def is_trainable(self):
        return False

    def predict(self, docs):
        scores = self.model(docs)
        return scores

    def set_annnotations(self, docs, scores):
        for doc in docs:
            for token in doc:
                token.tag_ = scores[some_index]

And then your config would have something like this:

[initialize.components.tagger]
path = "/path/to/models_pt"

The local path is only need on initialization and then after that it's saved in the wrapped thinc format as part of the pipeline. You'd use spacy assemble to set up the pipeline from the config.

This would be a useful kind of demo to add to one of our demo projects, too, so don't hesitate about asking more questions. (I still find the API side of things frustrating. It feels like we need SerializablePipe or something alongside TrainablePipe.)

[kanayer]

Thank you for your detailed explanation! Could you please explain more on what should go into convert_inputs= and convert_outputs=? The documentation on the Thinc website says convert_outputs returns a tuple of the PyTorch outputs. But what if my tagger model consists of 2 stages?

• The first stage is the training stage, where the model.pt is created and saved.
• The second stage is tagging stage. Where the model predictions are taken and applied to the data samples from the test dataset.

Am I correct to assume that the code from the tagging stage should be written into the predict and set_annotations sections of the custom component code? But then predict accepts docs as the input, would that be a problem?

To give you a better idea of how tagging works in my model, I'll attach the piece of code below:

This piece of code loads the trained model

checkpoint = torch.load(config['model_name'], map_location=device)
model.load_state_dict(checkpoint['model'])
logger.info(model)
model.eval()

This piece of code loads the data from the input_file and attaches them to examples and creates iterator

with io.open(args.input_file, 'r', encoding='utf-8') as f:
    sents = f.readlines()

Default = torchtext.data.RawField()
named_fields = [('word', WORD), ('raw', Default)]
examples = []
for sent in sents:
    syllables = [list(eojeol) for eojeol in sent.split()]
    print(syllables)
    examples.append(POSExample.fromsent(syllables, named_fields))

text_dataset = Dataset(examples, named_fields)
text_iter = torchtext.data.BucketIterator(text_dataset, batch_size=5, shuffle=False)

This piece of code infers the model and tags the data based on the predictions of the model

outputs = []
with torch.no_grad():
    for t in tqdm(text_iter):
        decoder_outputs, tagger_outputs, others = model.infer(t.word[0].to(device))
        for n in range(len(others['length'])):
            length = others['length'][n]

            tgt_id_seq = [others['sequence'][di][n].item() for di in range(length)]
            tag_id_seq = [tagger_outputs[n][di].item() for di in range(length)]

            result = []
            analyze = []
            for i, (tgt, pos) in enumerate(zip(tgt_id_seq, tag_id_seq)):
                if tgt == LEX.vocab.stoi[LEX.unk_token] and config['replace_unk']:
                    _, idx = others['attention_score'][i][n].max(1)
                    attn_idx = idx.item() - 1
                    if attn_idx < 0 or attn_idx >= len(t.raw[n]):
                        continue

                    result.append(t.raw[n][attn_idx] + "/" + POS_TAG.vocab.itos[pos])
                else:
                    result.append(LEX.vocab.itos[tgt] + "/" + POS_TAG.vocab.itos[pos])

            outputs.append(result)

The final predictions stored in outputs are then run through the function called syllable_to_eojeol and you would have morphemes along with their pos tags. For example,[('요즘', 'NNG'), ('처럼', 'JKB')].

Is it possible to write a custom pipeline component in this case? Do these parts from the tagging code have to be integrated into predict and set_annotations sections of the custom component code?

[adrianeboyd]

Yes, the tagging stage should go into predict and set_annotations. Since you're not training, it would probably be fine if everything is in set_annotations, but you'd still need to override predict so that it does nothing because nlp and nlp.pipe call predict+set_annotations. But if it's not hard to separate it into predict + set_annotations, then that would fit the API better overall.

Like I said, we don't have enough examples of this. There's the spacy-transformers example for the input/output conversion here:

https://github.com/explosion/spacy-transformers/blob/3f921f2832f235af1a8a6d14cfda0939bd7b828f/spacy_transformers/layers/transformer_model.py#L208-L233

And the default options from thinc:

https://github.com/explosion/thinc/blob/f2af6d2dc0de4196f243616a3e37d03db369f0fc/thinc/layers/pytorchwrapper.py#L149-L219

The main thing is to do the numpy/cupy<->torch conversions of the data that you're passing to the model.

[kanayer]

Oh, I see, thank you! I have a couple more questions.

1. Do I have to write functions def _convert_lstm_inputs and def _convert_lstm_outputs as in the transformer examples for my custom component? Or should it be handled all inside set_annotations?
2. I wrote the factory code for the custom component but I am not sure what exactly needs to be passed in the @Language.factory? Only nlp and name? Also, from where do functions like def make_torch_tagger take the inputs? From the config file or the CLI?

@Language.factory("torch_tagger")
def make_torch_tagger(nlp, name):
    return TorchPOSTagger(nlp.vocab, name)

3. Where do you pass the docs? I mean if I don't initialize it in the def initialize, but call in in set_annotations, where will it take it from?
4. Do I put only

[initialize.components.tagger]
path = "/path/to/models_pt"

this into the config file? Because when I call the config file only with this, I get the following error:

initialize -> components -> tagger Section 'initialize' is not defined

I tried to fill the file with [nlp] and [initialize] and [initialize_components] in the following way:

[nlp]
lang = "ko"
pipeline = []

[initialize]

[initialize_components]
factory = "torch_tagger"

[initialize.components.tagger]
path = "./model/models_pt"

But I get the error

initialize -> components -> tagger Section 'components' is not defined

Could you please help me understand what am I doing wrong?

[adrianeboyd]

You need placeholders for all the config sections, so (also . not _!):

[initialize]

[initialize.components]

[initialize.components.tagger]
path = "./model/models_pt"

The factory doesn't go in initialize, but in the separate [components.tagger] section. The settings in [components] are used every time the model is loaded, but [initialize] is only used once, either with spacy train or spacy assemble.

I suspect your case will work with the default functions convert_pytorch_default_inputs and convert_pytorch_default_outputs defined for PyTorchWrapper and you don't need to write custom ones. You'd mainly need to do the conversion in set_annotations from the Doc annotation to the right input for your model.