
import warnings
from pathlib import Path
from typing import List, Tuple, Union

import fire
from torch import nn

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer, PreTrainedModel
from transformers.utils import logging


lowerCamelCase      =logging.get_logger(__name__)


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		):
      UpperCamelCase__    :   List[str]					=				nn.ModuleList([src_layers[i] for i in layers_to_copy]		)
      assert len(UpperCamelCase__		) == len(UpperCamelCase__		), f'''{len(UpperCamelCase__		)} != {len(UpperCamelCase__		)}'''
      dest_layers.load_state_dict(layers_to_copy.state_dict()		)


lowerCamelCase      ={
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    1_2: {
        1: [0],  # This says that if the teacher has 12 layers and the student has 1, copy layer 0 of the teacher
        2: [0, 6],
        3: [0, 6, 1_1],
        4: [0, 4, 8, 1_1],
        6: [0, 2, 4, 7, 9, 1_1],
        9: [0, 1, 2, 4, 5, 7, 9, 1_0, 1_1],
        1_2: list(range(1_2)),
    },
    1_6: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 1_5],
        3: [0, 8, 1_5],
        4: [0, 5, 1_0, 1_5],
        6: [0, 3, 6, 9, 1_2, 1_5],
        8: [0, 2, 4, 6, 8, 1_0, 1_2, 1_5],
        9: [0, 1, 3, 5, 7, 9, 1_1, 1_3, 1_5],
        1_2: [0, 1, 2, 3, 4, 5, 6, 7, 9, 1_1, 1_3, 1_5],
        1_6: list(range(1_6)),
    },
    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
}
lowerCamelCase      ={
    # maps  num layers in student -> which teacher layers to copy.
    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
    1_2: {1: [1_1], 2: [5, 1_1], 3: [3, 7, 1_1], 6: [1, 3, 5, 8, 1_0, 1_1]},
    1_6: {1: [1_5], 4: [4, 9, 1_2, 1_5], 8: [1, 3, 5, 7, 9, 1_1, 1_3, 1_5]},
}


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      try:
            UpperCamelCase__    :   Dict					=				LAYERS_TO_COPY[n_teacher][n_student]
            return val
      except KeyError:
            if n_student != n_teacher:
                  warnings.warn(
                      f'''no hardcoded layers to copy for teacher {n_teacher} -> student {n_student}, defaulting to first'''
                      f''' {n_student}'''		)
            return list(range(UpperCamelCase__		)		)


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      if n_student > n_teacher:
            raise ValueError(f'''Cannot perform intermediate supervision for student {n_student} > teacher {n_teacher}'''		)
      elif n_teacher == n_student:
            return list(range(UpperCamelCase__		)		)
      elif n_student == 1:
            return [n_teacher - 1]
      else:
            return LAYERS_TO_SUPERVISE[n_teacher][n_student]


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ = "student" ,		UpperCamelCase__ = None ,		UpperCamelCase__ = None ,		UpperCamelCase__=False ,		UpperCamelCase__=None ,		UpperCamelCase__=None ,		**UpperCamelCase__ ,		):
      UpperCamelCase__    :   Optional[int]					=				'''encoder_layers and decoder_layers cannot be both None-- you would just have an identical teacher.'''
      assert (e is not None) or (d is not None), _msg
      if isinstance(UpperCamelCase__ ,		UpperCamelCase__		):
            AutoTokenizer.from_pretrained(UpperCamelCase__		).save_pretrained(UpperCamelCase__		)  # purely for convenience
            UpperCamelCase__    :   List[str]					=				AutoModelForSeqaSeqLM.from_pretrained(UpperCamelCase__		).eval()
      else:
            assert isinstance(UpperCamelCase__ ,		UpperCamelCase__		), f'''teacher must be a model or string got type {type(UpperCamelCase__		)}'''
      UpperCamelCase__    :   Union[str, Any]					=				teacher.config.to_diff_dict()

      try:
            UpperCamelCase__    :   str					=				teacher.config.encoder_layers, teacher.config.decoder_layers
            if e is None:
                  UpperCamelCase__    :   Optional[int]					=				teacher_e
            if d is None:
                  UpperCamelCase__    :   Optional[int]					=				teacher_d
            init_kwargs.update({'''encoder_layers''': e, '''decoder_layers''': d}		)
      except AttributeError:  # T5
            if hasattr(teacher.config ,		'''num_encoder_layers'''		):
                  UpperCamelCase__    :   Optional[int]					=				teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
            else:
                  UpperCamelCase__    :   List[str]					=				teacher.config.num_layers, teacher.config.num_decoder_layers
            if e is None:
                  UpperCamelCase__    :   Tuple					=				teacher_e
            if d is None:
                  UpperCamelCase__    :   Optional[int]					=				teacher_d
            if hasattr(teacher.config ,		'''num_encoder_layers'''		):
                  init_kwargs.update({'''num_encoder_layers''': e, '''num_decoder_layers''': d}		)
            else:
                  init_kwargs.update({'''num_layers''': e, '''num_decoder_layers''': d}		)

    # Kwargs to instantiate student: teacher kwargs with updated layer numbers + **extra_config_kwargs
      init_kwargs.update(UpperCamelCase__		)

      # Copy weights
      UpperCamelCase__    :   List[Any]					=				teacher.config_class(**UpperCamelCase__		)
      UpperCamelCase__    :   Dict					=				AutoModelForSeqaSeqLM.from_config(UpperCamelCase__		)
      # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
      UpperCamelCase__    :   int					=				student.load_state_dict(teacher.state_dict() ,		strict=UpperCamelCase__		)
      assert info.missing_keys == [], info.missing_keys  # every student key should have a teacher keys.

      if copy_first_teacher_layers:  # Our copying is done. We just log and save
            UpperCamelCase__    :   str					=				list(range(UpperCamelCase__		)		), list(range(UpperCamelCase__		)		)
            logger.info(
                f'''Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to'''
                f''' {save_path}'''		)
            student.save_pretrained(UpperCamelCase__		)
            return student, e_layers_to_copy, d_layers_to_copy

      # Decide which layers of the teacher to copy. Not exactly alternating -- we try to keep first and last layer.
      if e_layers_to_copy is None:
            UpperCamelCase__    :   List[int]					=				pick_layers_to_copy(UpperCamelCase__ ,		UpperCamelCase__		)
      if d_layers_to_copy is None:
            UpperCamelCase__    :   List[int]					=				pick_layers_to_copy(UpperCamelCase__ ,		UpperCamelCase__		)

      try:
            if hasattr(
                UpperCamelCase__ ,		'''prophetnet'''		):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
                  copy_layers(teacher.prophetnet.encoder.layers ,		student.prophetnet.encoder.layers ,		UpperCamelCase__		)
                  copy_layers(teacher.prophetnet.decoder.layers ,		student.prophetnet.decoder.layers ,		UpperCamelCase__		)
            else:
                  copy_layers(teacher.model.encoder.layers ,		student.model.encoder.layers ,		UpperCamelCase__		)
                  copy_layers(teacher.model.decoder.layers ,		student.model.decoder.layers ,		UpperCamelCase__		)
      except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
            copy_layers(teacher.encoder.block ,		student.encoder.block ,		UpperCamelCase__		)
            copy_layers(teacher.decoder.block ,		student.decoder.block ,		UpperCamelCase__		)
      logger.info(
          f'''Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}'''		)
      UpperCamelCase__    :   Tuple					=				{
          '''teacher_type''': teacher.config.model_type,
          '''copied_encoder_layers''': e_layers_to_copy,
          '''copied_decoder_layers''': d_layers_to_copy,
      }
      student.save_pretrained(UpperCamelCase__		)
      # Save information about copying for easier reproducibility

      return student, e_layers_to_copy, d_layers_to_copy


if __name__ == "__main__":
    fire.Fire(create_student_by_copying_alternating_layers)




import argparse
import json
from typing import List

from ltp import LTP

from transformers.models.bert.tokenization_bert import BertTokenizer


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # This defines a "chinese character" as anything in the CJK Unicode block:
      #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
      #
      # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
      # despite its name. The modern Korean Hangul alphabet is a different block,
      # as is Japanese Hiragana and Katakana. Those alphabets are used to write
      # space-separated words, so they are not treated specially and handled
      # like the all of the other languages.
      if (
          (cp >= 0x4e_00 and cp <= 0x9f_ff)
          or (cp >= 0x34_00 and cp <= 0x4d_bf)  #
          or (cp >= 0x2_00_00 and cp <= 0x2_a6_df)  #
          or (cp >= 0x2_a7_00 and cp <= 0x2_b7_3f)  #
          or (cp >= 0x2_b7_40 and cp <= 0x2_b8_1f)  #
          or (cp >= 0x2_b8_20 and cp <= 0x2_ce_af)  #
          or (cp >= 0xf9_00 and cp <= 0xfa_ff)
          or (cp >= 0x2_f8_00 and cp <= 0x2_fa_1f)  #
      ):  #
            return True

      return False


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # word like '180' or '身高' or '神'
      for char in word:
            UpperCamelCase__    :   Optional[Any]					=				ord(UpperCamelCase__		)
            if not _is_chinese_char(UpperCamelCase__		):
                  return 0
      return 1


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      UpperCamelCase__    :   Dict					=				set()

      for token in tokens:
            UpperCamelCase__    :   Optional[int]					=				len(UpperCamelCase__		) > 1 and is_chinese(UpperCamelCase__		)
            if chinese_word:
                  word_set.add(UpperCamelCase__		)
      UpperCamelCase__    :   List[Any]					=				list(UpperCamelCase__		)
      return word_list


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      if not chinese_word_set:
            return bert_tokens
      UpperCamelCase__    :   str					=				max([len(UpperCamelCase__		) for w in chinese_word_set]		)

      UpperCamelCase__    :   Optional[int]					=				bert_tokens
      UpperCamelCase__				,UpperCamelCase__    :   Tuple					=				0, len(UpperCamelCase__		)
      while start < end:
            UpperCamelCase__    :   Any					=				True
            if is_chinese(bert_word[start]		):
                  UpperCamelCase__    :   Any					=				min(end - start ,		UpperCamelCase__		)
                  for i in range(UpperCamelCase__ ,		1 ,		-1		):
                        UpperCamelCase__    :   List[Any]					=				''''''.join(bert_word[start : start + i]		)
                        if whole_word in chinese_word_set:
                              for j in range(start + 1 ,		start + i		):
                                    UpperCamelCase__    :   Optional[Any]					=				'''##''' + bert_word[j]
                              UpperCamelCase__    :   Dict					=				start + i
                              UpperCamelCase__    :   List[str]					=				False
                              break
            if single_word:
                  start += 1
      return bert_word


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		):
      UpperCamelCase__    :   str					=				[]

      for i in range(0 ,		len(UpperCamelCase__		) ,		1_0_0		):
            UpperCamelCase__    :   Optional[Any]					=				ltp_tokenizer.pipeline(lines[i : i + 1_0_0] ,		tasks=['''cws''']		).cws
            UpperCamelCase__    :   List[str]					=				[get_chinese_word(UpperCamelCase__		) for r in res]
            ltp_res.extend(UpperCamelCase__		)
      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      UpperCamelCase__    :   str					=				[]
      for i in range(0 ,		len(UpperCamelCase__		) ,		1_0_0		):
            UpperCamelCase__    :   Dict					=				bert_tokenizer(lines[i : i + 1_0_0] ,		add_special_tokens=UpperCamelCase__ ,		truncation=UpperCamelCase__ ,		max_length=5_1_2		)
            bert_res.extend(res['''input_ids''']		)
      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      UpperCamelCase__    :   Any					=				[]
      for input_ids, chinese_word in zip(UpperCamelCase__ ,		UpperCamelCase__		):
            UpperCamelCase__    :   Union[str, Any]					=				[]
            for id in input_ids:
                  UpperCamelCase__    :   List[str]					=				bert_tokenizer._convert_id_to_token(UpperCamelCase__		)
                  input_tokens.append(UpperCamelCase__		)
            UpperCamelCase__    :   Optional[Any]					=				add_sub_symbol(UpperCamelCase__ ,		UpperCamelCase__		)
            UpperCamelCase__    :   Any					=				[]
            # We only save pos of chinese subwords start with ##, which mean is part of a whole word.
            for i, token in enumerate(UpperCamelCase__		):
                  if token[:2] == "##":
                        UpperCamelCase__    :   Union[str, Any]					=				token[2:]
                        # save chinese tokens' pos
                        if len(UpperCamelCase__		) == 1 and _is_chinese_char(ord(UpperCamelCase__		)		):
                              ref_id.append(UpperCamelCase__		)
            ref_ids.append(UpperCamelCase__		)

      assert len(UpperCamelCase__		) == len(UpperCamelCase__		)

      return ref_ids


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      # For Chinese (Ro)Bert, the best result is from : RoBERTa-wwm-ext (https://github.com/ymcui/Chinese-BERT-wwm)
      # If we want to fine-tune these model, we have to use same tokenizer : LTP (https://github.com/HIT-SCIR/ltp)
      with open(args.file_name ,		'''r''' ,		encoding='''utf-8'''		) as f:
            UpperCamelCase__    :   List[Any]					=				f.readlines()
      UpperCamelCase__    :   Tuple					=				[line.strip() for line in data if len(UpperCamelCase__		) > 0 and not line.isspace()]  # avoid delimiter like '\u2029'
      UpperCamelCase__    :   int					=				LTP(args.ltp		)  # faster in GPU device
      UpperCamelCase__    :   int					=				BertTokenizer.from_pretrained(args.bert		)

      UpperCamelCase__    :   Dict					=				prepare_ref(UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		)

      with open(args.save_path ,		'''w''' ,		encoding='''utf-8'''		) as f:
            UpperCamelCase__    :   Tuple					=				[json.dumps(UpperCamelCase__		) + '''\n''' for ref in ref_ids]
            f.writelines(UpperCamelCase__		)


if __name__ == "__main__":
    lowerCamelCase      =argparse.ArgumentParser(description="prepare_chinese_ref")
    parser.add_argument(
        "--file_name",
        required=False,
        type=str,
        default="./resources/chinese-demo.txt",
        help="file need process, same as training data in lm",
    )
    parser.add_argument(
        "--ltp",
        required=False,
        type=str,
        default="./resources/ltp",
        help="resources for LTP tokenizer, usually a path",
    )
    parser.add_argument(
        "--bert",
        required=False,
        type=str,
        default="./resources/robert",
        help="resources for Bert tokenizer",
    )
    parser.add_argument(
        "--save_path",
        required=False,
        type=str,
        default="./resources/ref.txt",
        help="path to save res",
    )

    lowerCamelCase      =parser.parse_args()
    main(args)


