text: """
  .. todo::
      doc
  """
  
  __all__ = [
      "ElmoEmbedding"
  ]
  
  import codecs
  import json
  import os
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from .contextual_embedding import ContextualEmbedding
  from ..core import logger
  from ..core.vocabulary import Vocabulary
  from ..io.file_utils import cached_path, _get_embedding_url, PRETRAINED_ELMO_MODEL_DIR
  from ..modules.encoder._elmo import ElmobiLm, ConvTokenEmbedder
  
  
  class ElmoEmbedding(ContextualEmbedding):
      """
      使用ELMo的embedding。初始化之后，只需要传入words就可以得到对应的embedding。
      当前支持的使用名称初始化的模型:
          en: 即en-medium hidden_size 1024; output_size 12
          en-medium: hidden_size 2048; output_size 256
          en-origial: hidden_size 4096; output_size 512
          en-original-5.5b: hidden_size 4096; output_size 512
          en-small: hidden_size 1024; output_size 128
  
      Example::
      
          >>> import torch
          >>> from fastNLP import Vocabulary
          >>> from fastNLP.embeddings import ElmoEmbedding
          >>> vocab = Vocabulary().add_word_lst("The whether is good .".split())
          >>> # 使用不同层的concat的结果
          >>> embed = ElmoEmbedding(vocab, model_dir_or_name='en', layers='1,2', requires_grad=False)
          >>> words = torch.LongTensor([[vocab.to_index(word) for word in "The whether is good .".split()]])
          >>> outputs = embed(words)
          >>> outputs.size()
          >>> # torch.Size([1, 5, 2048])
  
          >>> # 使用不同层的weighted sum。
          >>> embed = ElmoEmbedding(vocab, model_dir_or_name='en', layers='mix', requires_grad=False)
          >>> embed.set_mix_weights_requires_grad()  # 使得weighted的权重是可以学习的，但ELMO的LSTM部分是不更新
  
      """
      
      def __init__(self, vocab: Vocabulary, model_dir_or_name: str = 'en', layers: str = '2', requires_grad: bool = True,
                   word_dropout=0.0, dropout=0.0, cache_word_reprs: bool = False):
          """
          
          :param vocab: 词表
          :param model_dir_or_name: 可以有两种方式调用预训练好的ELMo embedding：第一种是传入ELMo所在文件夹，该文件夹下面应该有两个文件，
              其中一个是以json为后缀的配置文件，另一个是以pkl为后缀的权重文件；第二种是传入ELMo版本的名称，将自动查看缓存中是否存在该模型，
              没有的话将自动下载并缓存。
          :param layers: str, 指定返回的层数(从0开始), 以,隔开不同的层。如果要返回第二层的结果'2', 返回后两层的结果'1,2'。不同的层的结果
              按照这个顺序concat起来，默认为'2'。'mix'会使用可学习的权重结合不同层的表示(权重是否可训练与requires_grad保持一致，
              初始化权重对三层结果进行mean-pooling, 可以通过ElmoEmbedding.set_mix_weights_requires_grad()方法只将mix weights设置为可学习。)
          :param requires_grad: bool, 该层是否需要gradient, 默认为False.
          :param float word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
          :param float dropout: 以多大的概率对embedding的表示进行Dropout。0.1即随机将10%的值置为0。
          :param cache_word_reprs: 可以选择对word的表示进行cache; 设置为True的话，将在初始化的时候为每个word生成对应的embedding，
              并删除character encoder，之后将直接使用cache的embedding。默认为False。
          """
          super(ElmoEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
          
          # 根据model_dir_or_name检查是否存在并下载
          if model_dir_or_name.lower() in PRETRAINED_ELMO_MODEL_DIR:
              model_url = _get_embedding_url('elmo', model_dir_or_name.lower())
              model_dir = cached_path(model_url, name='embedding')
              # 检查是否存在
          elif os.path.isdir(os.path.abspath(os.path.expanduser(model_dir_or_name))):
              model_dir = model_dir_or_name
          else:
              raise ValueError(f"Cannot recognize {model_dir_or_name}.")
          self.model = _ElmoModel(model_dir, vocab, cache_word_reprs=cache_word_reprs)
          num_layers = self.model.encoder.num_layers
          
          if layers == 'mix':
              self.layer_weights = nn.Parameter(torch.zeros(self.model.config['lstm']['n_layers'] + 1),
                                                requires_grad=requires_grad)
              self.gamma = nn.Parameter(torch.ones(1), requires_grad=requires_grad)
              self._get_outputs = self._get_mixed_outputs
              self._embed_size = self.model.config['lstm']['projection_dim'] * 2
          else:
              layers = list(map(int, layers.split(',')))
              assert len(layers) > 0, "Must choose at least one output, but got None."
              for layer in layers:
                  assert 0 <= layer <= num_layers, f"Layer index should be in range [0, {num_layers}], but got {layer}."
              self.layers = layers
              self._get_outputs = self._get_layer_outputs
              self._embed_size = len(self.layers) * self.model.config['lstm']['projection_dim'] * 2
          
          self.requires_grad = requires_grad
      
      def _get_mixed_outputs(self, outputs):
          # outputs: num_layers x batch_size x max_len x hidden_size
          # return: batch_size x max_len x hidden_size
          weights = F.softmax(self.layer_weights + 1 / len(outputs), dim=0).to(outputs)
          outputs = torch.einsum('l,lbij->bij', weights, outputs)
          return self.gamma.to(outputs) * outputs
      
      def set_mix_weights_requires_grad(self, flag=True):
          """
          当初始化ElmoEmbedding时layers被设置为mix时，可以通过调用该方法设置mix weights是否可训练。如果layers不是mix，调用
          该方法没有用。
          
          :param bool flag: 混合不同层表示的结果是否可以训练。
          :return:
          """
          if hasattr(self, 'layer_weights'):
              self.layer_weights.requires_grad = flag
              self.gamma.requires_grad = flag
      
      def _get_layer_outputs(self, outputs):
          if len(self.layers) == 1:
              outputs = outputs[self.layers[0]]
          else:
              outputs = torch.cat(tuple([*outputs[self.layers]]), dim=-1)
          
          return outputs
      
      def forward(self, words: torch.LongTensor):
          """
          计算words的elmo embedding表示。根据elmo文章中介绍的ELMO实际上是有2L+1层结果，但是为了让结果比较容易拆分，token的
          被重复了一次，使得实际上layer=0的结果是[token_embedding;token_embedding], 而layer=1的结果是[forward_hiddens;
          backward_hiddens].
  
          :param words: batch_size x max_len
          :return: torch.FloatTensor. batch_size x max_len x (512*len(self.layers))
          """
          words = self.drop_word(words)
          outputs = self._get_sent_reprs(words)
          if outputs is not None:
              return self.dropout(outputs)
          outputs = self.model(words)
          outputs = self._get_outputs(outputs)
          return self.dropout(outputs)
      
      def _delete_model_weights(self):
          for name in ['layers', 'model', 'layer_weights', 'gamma']:
              if hasattr(self, name):
                  delattr(self, name)
  
  
  class _ElmoModel(nn.Module):
      """
      该Module是ElmoEmbedding中进行所有的heavy lifting的地方。做的工作，包括
          (1) 根据配置，加载模型;
          (2) 根据vocab，对模型中的embedding进行调整. 并将其正确初始化
          (3) 保存一个words与chars的对应转换，获取时自动进行相应的转换
          (4) 设计一个保存token的embedding，允许缓存word的表示。
  
      """
      
      def __init__(self, model_dir: str, vocab: Vocabulary = None, cache_word_reprs: bool = False):
          super(_ElmoModel, self).__init__()
          self.model_dir = model_dir
          dir = os.walk(self.model_dir)
          config_file = None
          weight_file = None
          config_count = 0
          weight_count = 0
          for path, dir_list, file_list in dir:
              for file_name in file_list:
                  if file_name.__contains__(".json"):
                      config_file = file_name
                      config_count += 1
                  elif file_name.__contains__(".pkl"):
                      weight_file = file_name
                      weight_count += 1
          if config_count > 1 or weight_count > 1:
              raise Exception(f"Multiple config files(*.json) or weight files(*.hdf5) detected in {model_dir}.")
          elif config_count == 0 or weight_count == 0:
              raise Exception(f"No config file or weight file found in {model_dir}")
          with open(os.path.join(model_dir, config_file), 'r') as config_f:
              config = json.load(config_f)
          self.weight_file = os.path.join(model_dir, weight_file)
          self.config = config
          
          OOV_TAG = '<oov>'
          PAD_TAG = '<pad>'
          BOS_TAG = '<bos>'
          EOS_TAG = '<eos>'
          BOW_TAG = '<bow>'
          EOW_TAG = '<eow>'
          
          # For the model trained with character-based word encoder.
          char_lexicon = {}
          with codecs.open(os.path.join(model_dir, 'char.dic'), 'r', encoding='utf-8') as fpi:
              for line in fpi:
                  tokens = line.strip().split('\t')
                  if len(tokens) == 1:
                      tokens.insert(0, '\u3000')
                  token, i = tokens
                  char_lexicon[token] = int(i)
          
          # 做一些sanity check
          for special_word in [PAD_TAG, OOV_TAG, BOW_TAG, EOW_TAG]:
              assert special_word in char_lexicon, f"{special_word} not found in char.dic."
          
          # 从vocab中构建char_vocab
          char_vocab = Vocabulary(unknown=OOV_TAG, padding=PAD_TAG)
          # 需要保证<bow>与<eow>在里面
          char_vocab.add_word_lst([BOW_TAG, EOW_TAG, BOS_TAG, EOS_TAG])
          
          for word, index in vocab:
              char_vocab.add_word_lst(list(word))
          
          self.bos_index, self.eos_index, self._pad_index = len(vocab), len(vocab) + 1, vocab.padding_idx
          # 根据char_lexicon调整, 多设置一位，是预留给word padding的(该位置的char表示为全0表示)
          char_emb_layer = nn.Embedding(len(char_vocab) + 1, int(config['char_cnn']['embedding']['dim']),
                                        padding_idx=len(char_vocab))
          
          # 读入预训练权重 这里的elmo_model 包含char_cnn和 lstm 的 state_dict
          elmo_model = torch.load(os.path.join(self.model_dir, weight_file), map_location='cpu')
          
          char_embed_weights = elmo_model["char_cnn"]['char_emb_layer.weight']
          
          found_char_count = 0
          for char, index in char_vocab:  # 调整character embedding
              if char in char_lexicon:
                  index_in_pre = char_lexicon.get(char)
                  found_char_count += 1
              else:
                  index_in_pre = char_lexicon[OOV_TAG]
              char_emb_layer.weight.data[index] = char_embed_weights[index_in_pre]
          
          logger.info(f"{found_char_count} out of {len(char_vocab)} characters were found in pretrained elmo embedding.")
          # 生成words到chars的映射
          max_chars = config['char_cnn']['max_characters_per_token']
          self.register_buffer('words_to_chars_embedding', torch.full((len(vocab) + 2, max_chars),
                                                                  fill_value=len(char_vocab),
                                                                  dtype=torch.long))
          for word, index in list(iter(vocab)) + [(BOS_TAG, len(vocab)), (EOS_TAG, len(vocab) + 1)]:
              if len(word) + 2 > max_chars:
                  word = word[:max_chars - 2]
              if index == self._pad_index:
                  continue
              elif word == BOS_TAG or word == EOS_TAG:
                  char_ids = [char_vocab.to_index(BOW_TAG)] + [char_vocab.to_index(word)] + [
                      char_vocab.to_index(EOW_TAG)]
                  char_ids += [char_vocab.to_index(PAD_TAG)] * (max_chars - len(char_ids))
              else:
                  char_ids = [char_vocab.to_index(BOW_TAG)] + [char_vocab.to_index(c) for c in word] + [
                      char_vocab.to_index(EOW_TAG)]
                  char_ids += [char_vocab.to_index(PAD_TAG)] * (max_chars - len(char_ids))
              self.words_to_chars_embedding[index] = torch.LongTensor(char_ids)
          
          self.char_vocab = char_vocab
          
          self.token_embedder = ConvTokenEmbedder(
              config, self.weight_file, None, char_emb_layer)
          elmo_model["char_cnn"]['char_emb_layer.weight'] = char_emb_layer.weight
          self.token_embedder.load_state_dict(elmo_model["char_cnn"])
          
          self.output_dim = config['lstm']['projection_dim']
          
          # lstm encoder
          self.encoder = ElmobiLm(config)
          self.encoder.load_state_dict(elmo_model["lstm"])
          
          if cache_word_reprs:
              if config['char_cnn']['embedding']['dim'] > 0:  # 只有在使用了chars的情况下有用
                  logger.info("Start to generate cache word representations.")
                  batch_size = 320
                  # bos eos
                  word_size = self.words_to_chars_embedding.size(0)
                  num_batches = word_size // batch_size + \
                                int(word_size % batch_size != 0)
                  
                  self.cached_word_embedding = nn.Embedding(word_size,
                                                            config['lstm']['projection_dim'])
                  with torch.no_grad():
                      for i in range(num_batches):
                          words = torch.arange(i * batch_size,
                                               min((i + 1) * batch_size, word_size)).long()
                          chars = self.words_to_chars_embedding[words].unsqueeze(1)  # batch_size x 1 x max_chars
                          word_reprs = self.token_embedder(words.unsqueeze(1),
                                                           chars).detach()  # batch_size x 1 x config['encoder']['projection_dim']
                          self.cached_word_embedding.weight.data[words] = word_reprs.squeeze(1)
                      
                      logger.info("Finish generating cached word representations. Going to delete the character encoder.")
                  del self.token_embedder, self.words_to_chars_embedding
              else:
                  logger.info("There is no need to cache word representations, since no character information is used.")
      
      def forward(self, words):
          """
  
          :param words: batch_size x max_len
          :return: num_layers x batch_size x max_len x hidden_size
          """
          # 扩展<bos>, <eos>
          batch_size, max_len = words.size()
          expanded_words = words.new_zeros(batch_size, max_len + 2)  # 因为pad一定为0，
          seq_len = words.ne(self._pad_index).sum(dim=-1)
          expanded_words[:, 1:-1] = words
          expanded_words[:, 0].fill_(self.bos_index)
          expanded_words[torch.arange(batch_size).to(words), seq_len + 1] = self.eos_index
          seq_len = seq_len + 2
          zero_tensor = expanded_words.new_zeros(expanded_words.shape)
          mask = (expanded_words == zero_tensor).unsqueeze(-1)
          if hasattr(self, 'cached_word_embedding'):
              token_embedding = self.cached_word_embedding(expanded_words)
          else:
              if hasattr(self, 'words_to_chars_embedding'):
                  chars = self.words_to_chars_embedding[expanded_words]
              else:
                  chars = None
              token_embedding = self.token_embedder(expanded_words, chars)  # batch_size x max_len x embed_dim
          
          encoder_output = self.encoder(token_embedding, seq_len)
          if encoder_output.size(2) < max_len + 2:
              num_layers, _, output_len, hidden_size = encoder_output.size()
              dummy_tensor = encoder_output.new_zeros(num_layers, batch_size,
                                                      max_len + 2 - output_len, hidden_size)
              encoder_output = torch.cat((encoder_output, dummy_tensor), 2)
          sz = encoder_output.size()  # 2, batch_size, max_len, hidden_size
          token_embedding = token_embedding.masked_fill(mask, 0)
          token_embedding = torch.cat((token_embedding, token_embedding), dim=2).view(1, sz[1], sz[2], sz[3])
          encoder_output = torch.cat((token_embedding, encoder_output), dim=0)
          
          # 删除<eos>, <bos>. 这里没有精确地删除，但应该也不会影响最后的结果了。
          encoder_output = encoder_output[:, :, 1:-1]
          return encoder_output
  