text: from transformers import (BertModel, XLNetModel, GPT2Model,
                            OpenAIGPTModel, TransfoXLModel, XLMModel,
                            RobertaModel, DistilBertModel,
                            CamembertModel, AlbertModel,
                            XLMRobertaModel, FlaubertModel)
  
  import torch
  
  bert_archs = {'bert-base-uncased', 'bert-large-uncased',
                'bert-base-uncased', 'bert-large-uncased',
                'bert-base-multilingual-uncased', 'bert-base-multilingual-cased',
                'bert-base-chinese', 'bert-base-german-cased',
                'bert-large-uncased-whole-word-masking',
                'bert-large-cased-whole-word-masking',
                'bert-large-uncased-whole-word-masking-finetuned-squad',
                'bert-large-cased-whole-word-masking-finetuned-squad',
                'bert-base-cased-finetuned-mrpc', 'bert-base-german-dbmdz-cased',
                'bert-base-german-dbmdz-uncased', 'cl-tohoku/bert-base-japanese',
                'cl-tohoku/bert-base-japanese-whole-word-masking',
                'cl-tohoku/bert-base-japanese-char',
                'cl-tohoku/bert-base-japanese-char-whole-word-masking',
                'TurkuNLP/bert-base-finnish-cased-v1',
                'TurkuNLP/bert-base-finnish-uncased-v1',
                'wietsedv/bert-base-dutch-cased'}
  xlnet_archs = {'xlnet-base-cased', 'xlnet-large-cased'}
  gpt2_archs = {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
  xlm_archs = {'xlm-mlm-en-2048', 'xlm-mlm-ende-1024', 'xlm-mlm-enfr-1024',
               'xlm-mlm-enro-1024', 'xlm-mlm-xnli15-1024',
               'xlm-mlm-tlm-xnli15-1024', 'xlm-clm-enfr-1024', 'xlm-clm-ende-1024',
               'xlm-mlm-17-1280', 'xlm-mlm-100-1280'}
  roberta_archs = {'roberta-base', 'roberta-large', 'roberta-large-mnli',
                   'distilroberta-base', 'roberta-base-openai-detector',
                   'roberta-large-openai-detector'}
  
  distilbert_archs = {'distilbert-base-uncased', 'distilbert-base-cased',
                      'distilbert-base-uncased-distilled-squad',
                      'distilbert-base-cased-distilled-squad', 'distilgpt2',
                      'distilbert-base-german-cased',
                      'distilbert-base-multilingual-cased'}
  
  albert_archs = {'albert-base-v1', 'albert-large-v1', 'albert-xlarge-v1',
                  'albert-xxlarge-v1', 'albert-base-v2', 'albert-large-v2',
                  'albert-xlarge-v2', 'albert-xxlarge-v2'}
  
  xlmroberta_archs = {'xlm-roberta-base', 'xlm-roberta-large'}
  
  flaubert_archs = {'flaubert/flaubert_small_cased',
                    'flaubert/flaubert_base_cased',
                    'flaubert/flaubert_base_uncased',
                    'flaubert/flaubert_large_cased'}
  
  
  class Embedding(torch.nn.Module):
      def __init__(self, arch, model, freeze, token_strategy, sentence_strategy):
          super(Embedding, self).__init__()
          self.model = model
          for p in self.model.parameters():
              p.requires_grad = not freeze
          self.embedding_size = self.model.config.hidden_size
          self.vocab_size = self.model.config.vocab_size
          self.max_length = self.model.config.max_position_embeddings
          self.arch = arch
          self.token_strategy = token_strategy
          self.sentence_strategy = sentence_strategy
  
      def extra_repr(self) -> str:
          s = 'Architecture={arch}'
          s += ', Embedding Size={embedding_size}'
          s += ', Vocabulary Size={vocab_size}'
          s += ', Maximum Input Length={max_length}'
          return s.format(**self.__dict__)
  
      def get_embeddings(self, hidden_states, output_token_embeddings):
          token_vectors = None
          if output_token_embeddings:
              token_embeddings = torch.stack(hidden_states, dim=0)
              token_embeddings = token_embeddings.permute(1, 2, 0, 3)
              if self.token_strategy is None:
                  token_vectors = torch.cat((token_embeddings[:, :, -1, :],
                                             token_embeddings[:, :, -2, :],
                                             token_embeddings[:, :, -3, :],
                                             token_embeddings[:, :, -4, :]),
                                            dim=2)
              else:
                  token_vectors = self.token_strategy(token_embeddings)
  
          if self.sentence_strategy is None:
              sentence_vector = torch.mean(hidden_states[-2], dim=1)
          else:
              sentence_vector = self.sentence_strategy(hidden_states)
  
          return (sentence_vector,) + (token_vectors,)
  
  
  class BertEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in bert_archs, 'Invalid Architecture'
          assert type(model) == BertModel, 'Invalid Model'
          super(BertEmbedding, self).__init__(arch, model, freeze,
                                              token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='bert-base-uncased', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in bert_archs, 'Invalid Architecture'
          return cls(arch,
                     BertModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class XLNetEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in xlnet_archs, 'Invalid Architecture'
          assert type(model) == XLNetModel, 'Invalid Model'
          super(XLNetEmbedding, self).__init__(arch, model, freeze,
                                               token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='xlnet-base-cased', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in xlnet_archs, 'Invalid Architecture'
          return cls(arch,
                     XLNetModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[1]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class GPT2Embedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in gpt2_archs, 'Invalid Architecture'
          assert type(model) == GPT2Model, 'Invalid Model'
          super(GPT2Embedding, self).__init__(arch, model, freeze,
                                              token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='gpt2', freeze=True, token_strategy=None,
                          sentence_strategy=None):
          assert arch in gpt2_archs, 'Invalid Architecture'
          return cls(arch,
                     GPT2Model.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class GPTEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch == 'openai-gpt', 'Invalid Architecture'
          assert type(model) == OpenAIGPTModel, 'Invalid Model'
          super(GPTEmbedding, self).__init__(arch, model, freeze, token_strategy,
                                             sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='openai-gpt', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch == 'openai-gpt', 'Invalid Architecture'
          return cls(arch,
                     OpenAIGPTModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[1]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class TransfoXLEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch == 'transfo-xl-wt103', 'Invalid Architecture'
          assert type(model) == TransfoXLModel, 'Invalid Model'
          super(TransfoXLEmbedding, self).__init__(arch, model, freeze,
                                                   token_strategy,
                                                   sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='gpt2', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch == 'transfo-xl-wt103', 'Invalid Architecture'
          return cls(arch,
                     TransfoXLModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, output_token_embeddings=True):
          outputs = self.model(input_ids=input, output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class XLMEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in xlm_archs, 'Invalid Architecture'
          assert type(model) == XLMModel, 'Invalid Model'
          super(XLMEmbedding, self).__init__(arch, model, freeze,
                                             token_strategy,
                                             sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='xlm-mlm-en-2048', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in xlm_archs, 'Invalid Architecture'
          return cls(arch,
                     XLMModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[1]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class RobertaEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in roberta_archs, 'Invalid Architecture'
          assert type(model) == RobertaModel, 'Invalid Model'
          super(RobertaEmbedding, self).__init__(arch, model, freeze,
                                                 token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='bert-base-uncased', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in roberta_archs, 'Invalid Architecture'
          return cls(arch,
                     RobertaModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class DistilBertEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in distilbert_archs, 'Invalid Architecture'
          assert type(model) == DistilBertModel, 'Invalid Model'
          super(DistilBertEmbedding, self).__init__(arch, model, freeze,
                                                    token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='distilbert-base-uncased', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in distilbert_archs, 'Invalid Architecture'
          return cls(arch,
                     DistilBertModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[1]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class CamembertEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch == 'camembert-base', 'Invalid Architecture'
          assert type(model) == CamembertModel, 'Invalid Model'
          super(CamembertEmbedding, self).__init__(arch, model, freeze,
                                                   token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='camembert-base', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch == 'camembert-base', 'Invalid Architecture'
          return cls(arch,
                     CamembertModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class AlbertEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in albert_archs, 'Invalid Architecture'
          assert type(model) == AlbertModel, 'Invalid Model'
          super(AlbertEmbedding, self).__init__(arch, model, freeze,
                                                token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='albert-base-v1', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in albert_archs, 'Invalid Architecture'
          return cls(arch,
                     AlbertModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class XLMRobertaEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in xlmroberta_archs, 'Invalid Architecture'
          assert type(model) == XLMRobertaModel, 'Invalid Model'
          super(XLMRobertaEmbedding, self).__init__(arch, model, freeze,
                                                    token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='xlm-roberta-base', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in xlmroberta_archs, 'Invalid Architecture'
          return cls(arch,
                     XLMRobertaModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[2]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  
  
  class FlaubertEmbedding(Embedding):
      def __init__(self, arch, model, freeze=True, token_strategy=None,
                   sentence_strategy=None) -> None:
          assert arch in flaubert_archs, 'Invalid Architecture'
          assert type(model) == FlaubertModel, 'Invalid Model'
          super(FlaubertEmbedding, self).__init__(arch, model, freeze,
                                                  token_strategy, sentence_strategy)
  
      @classmethod
      def from_pretrained(cls, arch='flaubert/flaubert_small_cased', freeze=True,
                          token_strategy=None, sentence_strategy=None):
          assert arch in flaubert_archs, 'Invalid Architecture'
          return cls(arch,
                     FlaubertModel.from_pretrained(arch),
                     freeze,
                     token_strategy,
                     sentence_strategy)
  
      def forward(self, input, mask=None, output_token_embeddings=True):
          outputs = self.model(input_ids=input, attention_mask=mask,
                               output_hidden_states=True)
          hidden_states = outputs[1]
          return self.get_embeddings(hidden_states, output_token_embeddings)
  