# AUTOGENERATED! DO NOT EDIT! File to edit: dev/00_core.ipynb (unless otherwise specified).

__all__ = ['CustomTransformerModel', 'TransformersBaseTokenizer', 'TransformersVocab']

# Cell
import torch
from transformers import PreTrainedModel, PreTrainedTokenizer
from fastai.text.transform import BaseTokenizer, Vocab

import numpy as np
from typing import List, Collection
import collections

# Cell
class CustomTransformerModel(torch.nn.Module):
    def __init__(self, transformer_model: PreTrainedModel):
        super(CustomTransformerModel, self).__init__()
        self.transformer = transformer_model

    def forward(self, input_ids):
        # Return only the logits from the transfomer
        logits = self.transformer(input_ids)[0]
        return logits

# Cell
class TransformersBaseTokenizer(BaseTokenizer):
    """Wrapper around PreTrainedTokenizer to be compatible with fast.ai"""

    def __init__(
        self, pretrained_tokenizer: PreTrainedTokenizer, model_type="bert", **kwargs
    ):
        self._pretrained_tokenizer = pretrained_tokenizer
        self.max_seq_len = pretrained_tokenizer.max_len
        self.model_type = model_type

    def __call__(self, *args, **kwargs):
        return self

    def tokenizer(self, t: str) -> List[str]:
        """Limits the maximum sequence length and add the spesial tokens"""
        CLS = self._pretrained_tokenizer.cls_token
        SEP = self._pretrained_tokenizer.sep_token
        if self.model_type in ["roberta"]:
            tokens = self._pretrained_tokenizer.tokenize(t, add_prefix_space=True)[
                : self.max_seq_len - 2
            ]
        else:
            tokens = self._pretrained_tokenizer.tokenize(t)[: self.max_seq_len - 2]
        return [CLS] + tokens + [SEP]

# Cell
class TransformersVocab(Vocab):
    def __init__(self, tokenizer: PreTrainedTokenizer):
        super(TransformersVocab, self).__init__(itos=[])
        self.tokenizer = tokenizer

    def numericalize(self, t: Collection[str]) -> List[int]:
        "Convert a list of tokens `t` to their ids."
        return self.tokenizer.convert_tokens_to_ids(t)
        # return self.tokenizer.encode(t)

    def textify(self, nums: Collection[int], sep=" ") -> List[str]:
        "Convert a list of `nums` to their tokens."
        nums = np.array(nums).tolist()
        return (
            sep.join(self.tokenizer.convert_ids_to_tokens(nums))
            if sep is not None
            else self.tokenizer.convert_ids_to_tokens(nums)
        )

    def __getstate__(self):
        return {"itos": self.itos, "tokenizer": self.tokenizer}

    def __setstate__(self, state: dict):
        self.itos = state["itos"]
        self.tokenizer = state["tokenizer"]
        self.stoi = collections.defaultdict(
            int, {v: k for k, v in enumerate(self.itos)}
        )