

"""simple docstring"""






import os
import pickle
import unittest

from transformers import AutoTokenizer
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.models.bert_japanese.tokenization_bert_japanese import (
    VOCAB_FILES_NAMES,
    BertJapaneseTokenizer,
    CharacterTokenizer,
    JumanppTokenizer,
    MecabTokenizer,
    SudachiTokenizer,
    WordpieceTokenizer,
)
from transformers.testing_utils import custom_tokenizers, require_jumanpp, require_sudachi

from ...test_tokenization_common import TokenizerTesterMixin







@custom_tokenizers
class 	__UpperCamelCase    (    a__							,							unittest.TestCase						):
  _UpperCAmelCase          =		BertJapaneseTokenizer
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		True

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().setUp()

      _lowerCAmelCase       :   Tuple  =    [
          '[UNK]',
          '[CLS]',
          '[SEP]',
          'こんにちは',
          'こん',
          'にちは',
          'ばんは',
          '##こん',
          '##にちは',
          '##ばんは',
          '世界',
          '##世界',
          '、',
          '##、',
          '。',
          '##。',
      ]

      _lowerCAmelCase       :   str  =    os.path.join(self.tmpdirname    ,VOCAB_FILES_NAMES['vocab_file']      )
      with open(self.vocab_file    ,'w'    ,encoding='utf-8'      ) as vocab_writer:
          vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]      )      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    'こんにちは、世界。 \nこんばんは、世界。'
      _lowerCAmelCase       :   List[str]  =    'こんにちは 、 世界 。 こんばんは 、 世界 。'
      return input_text, output_text

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase       :   List[Any]  =    self.get_input_output_texts(_A      )
      _lowerCAmelCase       :   List[str]  =    tokenizer.encode(_A    ,add_special_tokens=_A      )
      _lowerCAmelCase       :   Optional[Any]  =    tokenizer.decode(_A    ,clean_up_tokenization_spaces=_A      )
      return text, ids

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    self.tokenizer_class(self.vocab_file      )

      _lowerCAmelCase       :   Tuple  =    tokenizer.tokenize('こんにちは、世界。\nこんばんは、世界。'      )
      self.assertListEqual(_A    ,['こんにちは', '、', '世界', '。', 'こん', '##ばんは', '、', '世界', '。']      )
      self.assertListEqual(tokenizer.convert_tokens_to_ids(_A      )    ,[3, 12, 10, 14, 4, 9, 12, 10, 14]      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    self.tokenizer_class(self.vocab_file    ,word_tokenizer_type='mecab'      )
      self.assertIsNotNone(_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    'こんにちは、世界。\nこんばんは、世界。'
      _lowerCAmelCase       :   Union[str, Any]  =    tokenizer.tokenize(_A      )
      self.assertListEqual(_A    ,['こんにちは', '、', '世界', '。', 'こん', '##ばんは', '、', '世界', '。']      )
      self.assertListEqual(tokenizer.convert_tokens_to_ids(_A      )    ,[3, 12, 10, 14, 4, 9, 12, 10, 14]      )

      _lowerCAmelCase       :   List[str]  =    os.path.join(self.tmpdirname    ,'tokenizer.bin'      )
      with open(_A    ,'wb'      ) as handle:
          pickle.dump(_A    ,_A      )

      with open(_A    ,'rb'      ) as handle:
          _lowerCAmelCase       :   Optional[int]  =    pickle.load(_A      )

      _lowerCAmelCase       :   Optional[int]  =    tokenizer_new.tokenize(_A      )

      self.assertListEqual(_A    ,_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    MecabTokenizer(mecab_dic='ipadic'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップルストア', 'で', 'iPhone', '8', 'が', '発売', 'さ', 'れ', 'た', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      try:
          _lowerCAmelCase       :   Any  =    MecabTokenizer(mecab_dic='unidic_lite'      )
      except ModuleNotFoundError:
          return

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iPhone', '8', 'が', '発売', 'さ', 'れ', 'た', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      try:
          _lowerCAmelCase       :   int  =    MecabTokenizer(mecab_dic='unidic'      )
      except ModuleNotFoundError:
          return

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iPhone', '8', 'が', '発売', 'さ', 'れ', 'た', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    MecabTokenizer(do_lower_case=_A    ,mecab_dic='ipadic'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップルストア', 'で', 'iphone', '8', 'が', '発売', 'さ', 'れ', 'た', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      try:
          _lowerCAmelCase       :   List[str]  =    MecabTokenizer(
              do_lower_case=_A    ,normalize_text=_A    ,mecab_option='-d /usr/local/lib/mecab/dic/jumandic'      )
      except RuntimeError:
          # if dict doesn't exist in the system, previous code raises this error.
          return

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['ｱｯﾌﾟﾙストア', 'で', 'iPhone', '８', 'が', '発売', 'さ', 'れた', '\u3000', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    MecabTokenizer(normalize_text=_A    ,mecab_dic='ipadic'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['ｱｯﾌﾟﾙストア', 'で', 'iPhone', '８', 'が', '発売', 'さ', 'れ', 'た', '　', '。']    ,)

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    self.tokenizer_class(self.vocab_file    ,word_tokenizer_type='sudachi'      )
      self.assertIsNotNone(_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    'こんにちは、世界。\nこんばんは、世界。'
      _lowerCAmelCase       :   str  =    tokenizer.tokenize(_A      )
      self.assertListEqual(_A    ,['こんにちは', '、', '世界', '。', 'こん', '##ばんは', '、', '世界', '。']      )
      self.assertListEqual(tokenizer.convert_tokens_to_ids(_A      )    ,[3, 12, 10, 14, 4, 9, 12, 10, 14]      )

      _lowerCAmelCase       :   Optional[int]  =    os.path.join(self.tmpdirname    ,'tokenizer.bin'      )
      with open(_A    ,'wb'      ) as handle:
          pickle.dump(_A    ,_A      )

      with open(_A    ,'rb'      ) as handle:
          _lowerCAmelCase       :   Any  =    pickle.load(_A      )

      _lowerCAmelCase       :   Tuple  =    tokenizer_new.tokenize(_A      )

      self.assertListEqual(_A    ,_A      )

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    SudachiTokenizer(sudachi_dict_type='core'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,[' ', '\t', 'アップル', 'ストア', 'で', 'iPhone', '8', ' ', 'が', ' ', ' ', '\n ', '発売', 'さ', 'れ', 'た', ' ', '。', ' ', ' ']    ,)

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    SudachiTokenizer(sudachi_dict_type='core'    ,sudachi_split_mode='A'      )

      self.assertListEqual(tokenizer.tokenize('外国人参政権'      )    ,['外国', '人', '参政', '権']      )

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    SudachiTokenizer(sudachi_dict_type='core'    ,sudachi_split_mode='B'      )

      self.assertListEqual(tokenizer.tokenize('外国人参政権'      )    ,['外国人', '参政権']      )

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    SudachiTokenizer(sudachi_dict_type='core'    ,sudachi_split_mode='C'      )

      self.assertListEqual(tokenizer.tokenize('外国人参政権'      )    ,['外国人参政権']      )

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    SudachiTokenizer(do_lower_case=_A    ,sudachi_dict_type='core'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,[' ', '\t', 'アップル', 'ストア', 'で', 'iphone', '8', ' ', 'が', ' ', ' ', '\n ', '発売', 'さ', 'れ', 'た', ' ', '。', ' ', ' ']    ,)

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    SudachiTokenizer(normalize_text=_A    ,sudachi_dict_type='core'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,[' ', '\t', 'ｱｯﾌﾟﾙ', 'ストア', 'で', 'iPhone', '８', ' ', 'が', ' ', ' ', '\n ', '発売', 'さ', 'れ', 'た', '\u3000', '。', ' ', ' ']    ,)

  @require_sudachi
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    SudachiTokenizer(trim_whitespace=_A    ,sudachi_dict_type='core'      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iPhone', '8', 'が', '発売', 'さ', 'れ', 'た', '。']    ,)

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    self.tokenizer_class(self.vocab_file    ,word_tokenizer_type='jumanpp'      )
      self.assertIsNotNone(_A      )

      _lowerCAmelCase       :   Dict  =    'こんにちは、世界。\nこんばんは、世界。'
      _lowerCAmelCase       :   List[Any]  =    tokenizer.tokenize(_A      )
      self.assertListEqual(_A    ,['こんにちは', '、', '世界', '。', 'こん', '##ばんは', '、', '世界', '。']      )
      self.assertListEqual(tokenizer.convert_tokens_to_ids(_A      )    ,[3, 12, 10, 14, 4, 9, 12, 10, 14]      )

      _lowerCAmelCase       :   List[str]  =    os.path.join(self.tmpdirname    ,'tokenizer.bin'      )
      with open(_A    ,'wb'      ) as handle:
          pickle.dump(_A    ,_A      )

      with open(_A    ,'rb'      ) as handle:
          _lowerCAmelCase       :   str  =    pickle.load(_A      )

      _lowerCAmelCase       :   List[str]  =    tokenizer_new.tokenize(_A      )

      self.assertListEqual(_A    ,_A      )

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    JumanppTokenizer()

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iPhone', '8', '\u3000', 'が', '\u3000', '\u3000', '\u3000', '発売', 'さ', 'れた', '\u3000', '。']    ,)

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    JumanppTokenizer(do_lower_case=_A      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iphone', '8', '\u3000', 'が', '\u3000', '\u3000', '\u3000', '発売', 'さ', 'れた', '\u3000', '。']    ,)

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    JumanppTokenizer(normalize_text=_A      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['ｱ', 'ｯ', 'ﾌ', 'ﾟ', 'ﾙ', 'ストア', 'で', 'iPhone', '８', '\u3000', 'が', '\u3000', '\u3000', '\u3000', '発売', 'さ', 'れた', '\u3000', '。']    ,)

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    JumanppTokenizer(trim_whitespace=_A      )

      self.assertListEqual(
          tokenizer.tokenize(' \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  '      )    ,['アップル', 'ストア', 'で', 'iPhone', '8', 'が', '発売', 'さ', 'れた', '。']    ,)

  @require_jumanpp
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    JumanppTokenizer()

      self.assertListEqual(
          tokenizer.tokenize('ありがとうございますm(_ _)ｍ見つけるのが大変です。'      )    ,['ありがとう', 'ございます', 'm(_ _)m', '見つける', 'の', 'が', '大変です', '。']    ,)

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    ['[UNK]', '[CLS]', '[SEP]', 'こんにちは', 'こん', 'にちは', 'ばんは', '##こん', '##にちは', '##ばんは']

      _lowerCAmelCase       :   Any  =    {}
      for i, token in enumerate(_A      ):
          _lowerCAmelCase       :   int  =    i
      _lowerCAmelCase       :   Tuple  =    WordpieceTokenizer(vocab=_A    ,unk_token='[UNK]'      )

      self.assertListEqual(tokenizer.tokenize(''      )    ,[]      )

      self.assertListEqual(tokenizer.tokenize('こんにちは'      )    ,['こんにちは']      )

      self.assertListEqual(tokenizer.tokenize('こんばんは'      )    ,['こん', '##ばんは']      )

      self.assertListEqual(tokenizer.tokenize('こんばんは こんばんにちは こんにちは'      )    ,['こん', '##ばんは', '[UNK]', 'こんにちは']      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    BertJapaneseTokenizer.from_pretrained('nlp-waseda/roberta-base-japanese-with-auto-jumanpp'      )
      _lowerCAmelCase       :   int  =    tokenizer.subword_tokenizer

      _lowerCAmelCase       :   Optional[Any]  =    subword_tokenizer.tokenize('国境 の 長い トンネル を 抜ける と 雪国 であった 。'      )
      self.assertListEqual(_A    ,['▁国境', '▁の', '▁長い', '▁トンネル', '▁を', '▁抜ける', '▁と', '▁雪', '国', '▁であった', '▁。']      )

      _lowerCAmelCase       :   Union[str, Any]  =    subword_tokenizer.tokenize('こんばんは こんばん にち は こんにちは'      )
      self.assertListEqual(_A    ,['▁こん', 'ばん', 'は', '▁こん', 'ばん', '▁に', 'ち', '▁は', '▁こんにちは']      )







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    self.tokenizer_class.from_pretrained('cl-tohoku/bert-base-japanese'      )

      _lowerCAmelCase       :   List[str]  =    tokenizer.encode('ありがとう。'    ,add_special_tokens=_A      )
      _lowerCAmelCase       :   Optional[Any]  =    tokenizer.encode('どういたしまして。'    ,add_special_tokens=_A      )

      _lowerCAmelCase       :   Tuple  =    tokenizer.build_inputs_with_special_tokens(_A      )
      _lowerCAmelCase       :   Tuple  =    tokenizer.build_inputs_with_special_tokens(_A    ,_A      )

      # 2 is for "[CLS]", 3 is for "[SEP]"
      assert encoded_sentence == [2] + text + [3]
      assert encoded_pair == [2] + text + [3] + text_a + [3]







@custom_tokenizers
class 	__UpperCamelCase    (    a__							,							unittest.TestCase						):
  _UpperCAmelCase          =		BertJapaneseTokenizer
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().setUp()

      _lowerCAmelCase       :   Union[str, Any]  =    ['[UNK]', '[CLS]', '[SEP]', 'こ', 'ん', 'に', 'ち', 'は', 'ば', '世', '界', '、', '。']

      _lowerCAmelCase       :   Tuple  =    os.path.join(self.tmpdirname    ,VOCAB_FILES_NAMES['vocab_file']      )
      with open(self.vocab_file    ,'w'    ,encoding='utf-8'      ) as vocab_writer:
          vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]      )      )

  def __lowerCamelCase		(		self    ,**_A      ):
      '''simple docstring'''

      return BertJapaneseTokenizer.from_pretrained(self.tmpdirname    ,subword_tokenizer_type='character'    ,**_A      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    'こんにちは、世界。 \nこんばんは、世界。'
      _lowerCAmelCase       :   int  =    'こ ん に ち は 、 世 界 。 こ ん ば ん は 、 世 界 。'
      return input_text, output_text

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass  # TODO add if relevant

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    self.tokenizer_class(self.vocab_file    ,subword_tokenizer_type='character'      )

      _lowerCAmelCase       :   Any  =    tokenizer.tokenize('こんにちは、世界。 \nこんばんは、世界。'      )
      self.assertListEqual(
          _A    ,['こ', 'ん', 'に', 'ち', 'は', '、', '世', '界', '。', 'こ', 'ん', 'ば', 'ん', 'は', '、', '世', '界', '。']      )
      self.assertListEqual(
          tokenizer.convert_tokens_to_ids(_A      )    ,[3, 4, 5, 6, 7, 11, 9, 10, 12, 3, 4, 8, 4, 7, 11, 9, 10, 12]      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    ['[UNK]', '[CLS]', '[SEP]', 'こ', 'ん', 'に', 'ち', 'は', 'ば', '世', '界', '、', '。']

      _lowerCAmelCase       :   Tuple  =    {}
      for i, token in enumerate(_A      ):
          _lowerCAmelCase       :   int  =    i
      _lowerCAmelCase       :   Dict  =    CharacterTokenizer(vocab=_A    ,unk_token='[UNK]'      )

      self.assertListEqual(tokenizer.tokenize(''      )    ,[]      )

      self.assertListEqual(tokenizer.tokenize('こんにちは'      )    ,['こ', 'ん', 'に', 'ち', 'は']      )

      self.assertListEqual(tokenizer.tokenize('こんにちほ'      )    ,['こ', 'ん', 'に', 'ち', '[UNK]']      )







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    self.tokenizer_class.from_pretrained('cl-tohoku/bert-base-japanese-char'      )

      _lowerCAmelCase       :   List[str]  =    tokenizer.encode('ありがとう。'    ,add_special_tokens=_A      )
      _lowerCAmelCase       :   Dict  =    tokenizer.encode('どういたしまして。'    ,add_special_tokens=_A      )

      _lowerCAmelCase       :   int  =    tokenizer.build_inputs_with_special_tokens(_A      )
      _lowerCAmelCase       :   Any  =    tokenizer.build_inputs_with_special_tokens(_A    ,_A      )

      # 2 is for "[CLS]", 3 is for "[SEP]"
      assert encoded_sentence == [2] + text + [3]
      assert encoded_pair == [2] + text + [3] + text_a + [3]







@custom_tokenizers
class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    'cl-tohoku/bert-base-japanese'
      _lowerCAmelCase       :   int  =    AutoTokenizer.from_pretrained(_A      )
      self.assertIsInstance(_A    ,_A      )







class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    'cl-tohoku/bert-base-japanese'
      with self.assertLogs('transformers'    ,level='WARNING'      ) as cm:
          BertTokenizer.from_pretrained(_A      )
          self.assertTrue(
              cm.records[0].message.startswith(
                  'The tokenizer class you load from this checkpoint is not the same type as the class this function'
                  ' is called from.'      )      )
      _lowerCAmelCase       :   Tuple  =    'bert-base-cased'
      with self.assertLogs('transformers'    ,level='WARNING'      ) as cm:
          BertJapaneseTokenizer.from_pretrained(_A      )
          self.assertTrue(
              cm.records[0].message.startswith(
                  'The tokenizer class you load from this checkpoint is not the same type as the class this function'
                  ' is called from.'      )      )








"""simple docstring"""






import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCAmelCase						=		logging.get_logger(__name__)


_lowerCAmelCase						=		"""https://openaipublic.azureedge.net/jukebox/models/"""
_lowerCAmelCase						=		{
    """jukebox-1b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """1b_lyrics/prior_level_2.pth.tar""",
    ],
    """jukebox-5b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """5b_lyrics/prior_level_2.pth.tar""",
    ],
}

def 		lowerCamelCase__		(					_lowerCamelCase   ):

    '''simple docstring'''



    if key.endswith('.model.1.bias'   ) and len(key.split('.'   )   ) > 10:
        _lowerCAmelCase       :   Optional[int]  =    key.replace('.model.1.bias' ,   '.conv1d_1.bias'   )
    elif key.endswith('.model.1.weight'   ) and len(key.split('.'   )   ) > 10:
        _lowerCAmelCase       :   Optional[int]  =    key.replace('.model.1.weight' ,   '.conv1d_1.weight'   )
    elif key.endswith('.model.3.bias'   ) and len(key.split('.'   )   ) > 10:
        _lowerCAmelCase       :   int  =    key.replace('.model.3.bias' ,   '.conv1d_2.bias'   )
    elif key.endswith('.model.3.weight'   ) and len(key.split('.'   )   ) > 10:
        _lowerCAmelCase       :   Tuple  =    key.replace('.model.3.weight' ,   '.conv1d_2.weight'   )

    if "conditioner_blocks.0." in key:
        _lowerCAmelCase       :   Dict  =    key.replace('conditioner_blocks.0' ,   'conditioner_blocks'   )

    if "prime_prior" in key:
        _lowerCAmelCase       :   str  =    key.replace('prime_prior' ,   'encoder'   )

    if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
        _lowerCAmelCase       :   Optional[Any]  =    key.replace('.emb.' ,   '.'   )

    if key.endswith('k'   ):  # replace vqvae.X.k with vqvae.X.codebook
        return key.replace('.k' ,   '.codebook'   )
    if "y_emb." in key:
        return key.replace('y_emb.' ,   'metadata_embedding.'   )

    if "x_emb.emb." in key:
        _lowerCAmelCase       :   Any  =    key.replace('0.x_emb.emb' ,   'embed_tokens'   )

    if "prime_state_ln" in key:
        return key.replace('prime_state_ln' ,   'encoder.final_layer_norm'   )
    if ".ln" in key:
        return key.replace('.ln' ,   '.layer_norm'   )
    if "_ln" in key:
        return key.replace('_ln' ,   '_layer_norm'   )

    if "prime_state_proj" in key:
        return key.replace('prime_state_proj' ,   'encoder.proj_in'   )
    if "prime_x_out" in key:
        return key.replace('prime_x_out' ,   'encoder.lm_head'   )
    if "prior.x_out" in key:
        return key.replace('x_out' ,   'fc_proj_out'   )
    if "x_emb" in key:
        return key.replace('x_emb' ,   'embed_tokens'   )

    return key

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Optional[int]  =    {}
    import re

    _lowerCAmelCase       :   Optional[Any]  =    re.compile(R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'   )
    _lowerCAmelCase       :   Optional[int]  =    re.compile(
        R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'   )
    _lowerCAmelCase       :   Dict  =    re.compile(R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'   )

    _lowerCAmelCase       :   Tuple  =    re.compile(R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'   )
    _lowerCAmelCase       :   Union[str, Any]  =    re.compile(
        R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'   )
    _lowerCAmelCase       :   Tuple  =    re.compile(R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'   )

    _lowerCAmelCase       :   Optional[int]  =    re.compile(R'conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)'   )
    _lowerCAmelCase       :   Dict  =    re.compile(
        R'conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'   )
    _lowerCAmelCase       :   List[str]  =    re.compile(R'conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)'   )

    for original_key, value in state_dict.items():
        # rename vqvae.encoder keys
        if re_encoder_block_conv_in.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   int  =    re_encoder_block_conv_in.match(_lowerCamelCase   )
            _lowerCAmelCase       :   int  =    regex_match.groups()
            _lowerCAmelCase       :   Any  =    int(groups[2]   ) * 2 + int(groups[3]   )
            _lowerCAmelCase       :   Dict  =    f"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"""
            _lowerCAmelCase       :   Optional[int]  =    re_encoder_block_conv_in.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_encoder_block_resnet.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Dict  =    re_encoder_block_resnet.match(_lowerCamelCase   )
            _lowerCAmelCase       :   Union[str, Any]  =    regex_match.groups()
            _lowerCAmelCase       :   Dict  =    int(groups[2]   ) * 2 + int(groups[3]   )
            _lowerCAmelCase       :   Tuple  =    {'1': 1, '3': 2}[groups[-2]]
            _lowerCAmelCase       :   Union[str, Any]  =    f"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."""
            _lowerCAmelCase       :   Optional[int]  =    f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
            _lowerCAmelCase       :   Optional[int]  =    prefix + resnet_block
            _lowerCAmelCase       :   Dict  =    re_encoder_block_resnet.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_encoder_block_proj_out.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   str  =    re_encoder_block_proj_out.match(_lowerCamelCase   )
            _lowerCAmelCase       :   Union[str, Any]  =    regex_match.groups()
            _lowerCAmelCase       :   Dict  =    f"""encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"""
            _lowerCAmelCase       :   Any  =    re_encoder_block_proj_out.sub(_lowerCamelCase ,   _lowerCamelCase   )

        # rename vqvae.decoder keys
        elif re_decoder_block_conv_out.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Optional[int]  =    re_decoder_block_conv_out.match(_lowerCamelCase   )
            _lowerCAmelCase       :   List[Any]  =    regex_match.groups()
            _lowerCAmelCase       :   Any  =    int(groups[2]   ) * 2 + int(groups[3]   ) - 2
            _lowerCAmelCase       :   Dict  =    f"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"""
            _lowerCAmelCase       :   Dict  =    re_decoder_block_conv_out.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_decoder_block_resnet.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Any  =    re_decoder_block_resnet.match(_lowerCamelCase   )
            _lowerCAmelCase       :   Dict  =    regex_match.groups()
            _lowerCAmelCase       :   Dict  =    int(groups[2]   ) * 2 + int(groups[3]   ) - 2
            _lowerCAmelCase       :   Dict  =    {'1': 1, '3': 2}[groups[-2]]
            _lowerCAmelCase       :   int  =    f"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."""
            _lowerCAmelCase       :   Optional[int]  =    f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
            _lowerCAmelCase       :   List[Any]  =    prefix + resnet_block
            _lowerCAmelCase       :   str  =    re_decoder_block_resnet.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_decoder_block_proj_in.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Any  =    re_decoder_block_proj_in.match(_lowerCamelCase   )
            _lowerCAmelCase       :   List[Any]  =    regex_match.groups()
            _lowerCAmelCase       :   str  =    f"""decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"""
            _lowerCAmelCase       :   str  =    re_decoder_block_proj_in.sub(_lowerCamelCase ,   _lowerCamelCase   )

        # rename prior cond.model to upsampler.upsample_block and resnet
        elif re_prior_cond_conv_out.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   List[Any]  =    re_prior_cond_conv_out.match(_lowerCamelCase   )
            _lowerCAmelCase       :   Any  =    regex_match.groups()
            _lowerCAmelCase       :   Optional[int]  =    int(groups[1]   ) * 2 + int(groups[2]   ) - 2
            _lowerCAmelCase       :   Any  =    f"""conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"""
            _lowerCAmelCase       :   List[str]  =    re_prior_cond_conv_out.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_prior_cond_resnet.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Dict  =    re_prior_cond_resnet.match(_lowerCamelCase   )
            _lowerCAmelCase       :   Tuple  =    regex_match.groups()
            _lowerCAmelCase       :   Any  =    int(groups[1]   ) * 2 + int(groups[2]   ) - 2
            _lowerCAmelCase       :   Tuple  =    {'1': 1, '3': 2}[groups[-2]]
            _lowerCAmelCase       :   List[Any]  =    f"""conditioner_blocks.upsampler.upsample_block.{block_index}."""
            _lowerCAmelCase       :   List[str]  =    f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
            _lowerCAmelCase       :   Dict  =    prefix + resnet_block
            _lowerCAmelCase       :   List[str]  =    re_prior_cond_resnet.sub(_lowerCamelCase ,   _lowerCamelCase   )

        elif re_prior_cond_proj_in.fullmatch(_lowerCamelCase   ):
            _lowerCAmelCase       :   Dict  =    re_prior_cond_proj_in.match(_lowerCamelCase   )
            _lowerCAmelCase       :   List[str]  =    regex_match.groups()
            _lowerCAmelCase       :   List[Any]  =    f"""conditioner_blocks.upsampler.proj_in.{groups[-1]}"""
            _lowerCAmelCase       :   Dict  =    re_prior_cond_proj_in.sub(_lowerCamelCase ,   _lowerCamelCase   )

        # keep original key
        else:
            _lowerCAmelCase       :   Optional[Any]  =    original_key

        _lowerCAmelCase       :   List[Any]  =    replace_key(_lowerCamelCase   )

        if f"""{key_prefix}.{key}""" not in model_state_dict or key is None:
            print(f"""failed converting {original_key} to {key}, does not match"""   )

        # handle missmatched shape
        elif value.shape != model_state_dict[f"""{key_prefix}.{key}"""].shape:
            _lowerCAmelCase       :   Dict  =    model_state_dict[f"""{key_prefix}.{key}"""]
            print(f"""{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"""   )
            _lowerCAmelCase       :   Optional[int]  =    original_key

        _lowerCAmelCase       :   Union[str, Any]  =    original_key
        _lowerCAmelCase       :   Optional[Any]  =    value

    return new_dict

@torch.no_grad()
def 		lowerCamelCase__		(					_lowerCamelCase=None ,   _lowerCamelCase=None   ):

    '''simple docstring'''



    for file in MODEL_MAPPING[model_name]:
        if not os.path.isfile(f"""{pytorch_dump_folder_path}/{file.split('/'   )[-1]}"""   ):
            _lowerCAmelCase       :   str  =    requests.get(f"""{PREFIX}{file}""" ,   allow_redirects=_lowerCamelCase   )
            os.makedirs(f"""{pytorch_dump_folder_path}/""" ,   exist_ok=_lowerCamelCase   )
            open(f"""{pytorch_dump_folder_path}/{file.split('/'   )[-1]}""" ,   'wb'   ).write(r.content   )

    _lowerCAmelCase       :   Union[str, Any]  =    MODEL_MAPPING[model_name.split('/'   )[-1]]

    _lowerCAmelCase       :   Optional[Any]  =    JukeboxConfig.from_pretrained(_lowerCamelCase   )
    _lowerCAmelCase       :   List[str]  =    JukeboxModel(_lowerCamelCase   )

    _lowerCAmelCase       :   int  =    []
    _lowerCAmelCase       :   Any  =    {}
    for i, dict_name in enumerate(_lowerCamelCase   ):
        _lowerCAmelCase       :   Union[str, Any]  =    torch.load(f"""{pytorch_dump_folder_path}/{dict_name.split('/'   )[-1]}"""   )['model']

        _lowerCAmelCase       :   Optional[Any]  =    {}
        for k in old_dic.keys():
            if k.endswith('.b'   ):
                _lowerCAmelCase       :   int  =    old_dic[k]
            elif k.endswith('.w'   ):
                _lowerCAmelCase       :   Tuple  =    old_dic[k]
            elif "level_2" not in dict_name and "cond.model." in k:
                _lowerCAmelCase       :   str  =    old_dic[k]
            else:
                _lowerCAmelCase       :   Optional[Any]  =    old_dic[k]

        _lowerCAmelCase       :   List[str]  =    'vqvae' if i == 0 else f"""priors.{3 - i}"""
        _lowerCAmelCase       :   Tuple  =    fix_jukebox_keys(_lowerCamelCase ,   model.state_dict() ,   _lowerCamelCase ,   _lowerCamelCase   )
        weight_dict.append(_lowerCamelCase   )

    _lowerCAmelCase       :   List[Any]  =    weight_dict.pop(0   )
    model.vqvae.load_state_dict(_lowerCamelCase   )
    for i in range(len(_lowerCamelCase   )   ):
        model.priors[i].load_state_dict(weight_dict[2 - i]   )

    Path(_lowerCamelCase   ).mkdir(exist_ok=_lowerCamelCase   )
    with open(f"""{pytorch_dump_folder_path}/mapping.json""" ,   'w'   ) as txtfile:
        json.dump(_lowerCamelCase ,   _lowerCamelCase   )

    print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""   )
    model.save_pretrained(_lowerCamelCase   )

    return weight_dict


if __name__ == "__main__":
 _lowerCAmelCase						=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--model_name""",
     default="""jukebox-5b-lyrics""",
     type=str,
     help="""Name of the model you'd like to convert.""",
 )
 parser.add_argument(
     """--pytorch_dump_folder_path""",
     default="""jukebox-5b-lyrics-converted""",
     type=str,
     help="""Path to the output PyTorch model directory.""",
 )
 _lowerCAmelCase						=		parser.parse_args()
 convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)





