text: # AUTOGENERATED! DO NOT EDIT! File to edit: 00_fastbert.ipynb (unless otherwise specified).
  
  __all__ = ['TransformersBaseTokenizer', 'SeqTokenizer', 'SeqTokenizeProcessor', 'segment', 'CustomTransformerModel',
             'TransformersVocab']
  
  # Cell
  # pytorch
  import torch
  
  # transformers
  from transformers import PreTrainedModel, PreTrainedTokenizer, PretrainedConfig
  from transformers import BertForSequenceClassification, BertTokenizer, BertConfig
  
  # fast.ai
  from fastai import *
  from fastai.text import *
  
  # Cell
  class TransformersBaseTokenizer(BaseTokenizer):
      """Wrapper around PreTrainedTokenizer to be compatible with fast.ai"""
      def __init__(self, pretrained_tokenizer: PreTrainedTokenizer, model_type = 'bert', max_len=64, **kwargs):
          self._pretrained_tokenizer = pretrained_tokenizer
          self.max_seq_len = max_len
          self.model_type = model_type
  
      def __call__(self, *args, **kwargs):
          return self
  
      def tokenizer(self, t:List[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
          turns = [self._pretrained_tokenizer.tokenize(turn) for turn in t]
          tokens = [CLS] + turns[0] + turns[1] + [SEP] + turns[2] + [SEP]
          return tokens[:self.max_seq_len]
  
  # Cell
  class SeqTokenizer(Tokenizer):
      "Put together rules and a tokenizer function to tokenize text with multiprocessing."
      def __init__(self, tok_func:Callable=SpacyTokenizer, lang:str='en', pre_rules:ListRules=None,
                   post_rules:ListRules=None, special_cases:Collection[str]=None, n_cpus:int=None):
          self.tok_func,self.lang,self.special_cases = tok_func,lang,special_cases
          self.pre_rules  = ifnone(pre_rules,  defaults.text_pre_rules )
          self.post_rules = ifnone(post_rules, defaults.text_post_rules)
          self.special_cases = special_cases if special_cases is not None else defaults.text_spec_tok
          self.n_cpus = ifnone(n_cpus, defaults.cpus)
  
      def process_text(self, t:List[str], tok:BaseTokenizer) -> List[str]:
          "Process one text `t` with tokenizer `tok`."
          for rule in self.pre_rules: t = rule(t)
          toks = tok.tokenizer(t)
          for rule in self.post_rules: toks = rule(toks)
          return toks
  
      def _process_all_1(self, texts:Collection[List[str]]) -> List[List[str]]:
          "Process a list of `texts` in one process."
          tok = self.tok_func(self.lang)
          if self.special_cases: tok.add_special_cases(self.special_cases)
          return [self.process_text(t, tok) for t in texts]
  
      def process_all(self, texts:Collection[List[str]]) -> List[List[str]]:
          "Process a list of `texts`."
          if self.n_cpus <= 1: return self._process_all_1(texts)
          with ProcessPoolExecutor(self.n_cpus) as e:
              return sum(e.map(self._process_all_1, partition_by_cores(texts, self.n_cpus)), [])
  
  # Cell
  class SeqTokenizeProcessor(TokenizeProcessor):
      "`PreProcessor` that tokenizes the texts in `ds`."
      def __init__(self, ds:ItemList=None, tokenizer:Tokenizer=None, chunksize:int=10000,
                   mark_fields:bool=False, include_bos:bool=True, include_eos:bool=False):
          self.tokenizer,self.chunksize,self.mark_fields = ifnone(tokenizer, Tokenizer()),chunksize,mark_fields
          self.include_bos, self.include_eos = include_bos, include_eos
  
      def process_one(self, item):
          return self.tokenizer._process_all_1(item)[0]
  
      def process(self, ds):
          tokens = []
          for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):
              tokens += self.tokenizer.process_all(ds.items[i:i+self.chunksize])
          ds.items = tokens
  
  # Cell
  def segment(input_ids):
      """util function for token_type_ids in bert"""
      segment_ids = input_ids.clone().cpu() # make sure VRAM will not explode
      segs = (segment_ids==102).nonzero().cpu().numpy()
      state = -1
      # current tensor
      cur = -1
      for seg in segs:
          if cur != seg[0]:
              cur = seg[0]
              segment_ids[cur][:seg[1]] = 0
              segment_ids[cur] = \
              (segment_ids[cur]!=0).type(segment_ids[cur].type())
  
      segs_set = set(segs.transpose()[0])
      if segs_set != segment_ids.shape[0]:
          for i in (set(range(segment_ids.shape[0])) - segs_set):
              segment_ids[i] = 0
  
      return segment_ids
  
  # Cell
  class CustomTransformerModel(nn.Module):
      """custom transformer model for fast.ai"""
      def __init__(self, transformer_model: PreTrainedModel):
          super(CustomTransformerModel,self).__init__()
          self.transformer = transformer_model
  
      def forward(self, input_ids):
  
          attention_mask = (input_ids!=0).type(input_ids.type())
          segmentation_mask = segment(input_ids).type(input_ids.type())
  
          logits = self.transformer(input_ids,
                                    attention_mask=attention_mask,
                                    token_type_ids=segmentation_mask)[0]
  
          return logits
  
  # 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)