


"""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 							A_      (       __snake_case     ,	unittest.TestCase ):
  lowerCAmelCase__      			=	BertJapaneseTokenizer
  lowerCAmelCase__      			=	False
  lowerCAmelCase__      			=	True


  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     super().setUp()

     _lowerCamelCase   :				Optional[Any]							=     [
         "[UNK]",
         "[CLS]",
         "[SEP]",
         "こんにちは",
         "こん",
         "にちは",
         "ばんは",
         "##こん",
         "##にちは",
         "##ばんは",
         "世界",
         "##世界",
         "、",
         "##、",
         "。",
         "##。",
     ]

     _lowerCamelCase   :				Dict							=     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 							_lowercase					(			self:						Any							,__lowerCAmelCase:						List[str]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Any							=     "こんにちは、世界。 \nこんばんは、世界。"
     _lowerCamelCase   :				List[Any]							=     "こんにちは 、 世界 。 こんばんは 、 世界 。"
     return input_text, output_text


  def 							_lowercase					(			self:						int							,__lowerCAmelCase:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     self.get_input_output_texts(A_    )
     _lowerCamelCase   :				Dict							=     tokenizer.encode(A_							,add_special_tokens=A_    )
     _lowerCamelCase   :				Union[str, Any]							=     tokenizer.decode(A_							,clean_up_tokenization_spaces=A_    )
     return text, ids


  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						str    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						Optional[int]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     self.tokenizer_class(self.vocab_file    )

     _lowerCamelCase   :				Optional[int]							=     tokenizer.tokenize("こんにちは、世界。\nこんばんは、世界。"    )
     self.assertListEqual(A_							,["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]    )
     self.assertListEqual(tokenizer.convert_tokens_to_ids(A_    )							,[3, 12, 10, 14, 4, 9, 12, 10, 14]    )


  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     self.tokenizer_class(self.vocab_file							,word_tokenizer_type="mecab"    )
     self.assertIsNotNone(A_    )

     _lowerCamelCase   :				List[str]							=     "こんにちは、世界。\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   :				Union[str, Any]							=     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_    )


  def 							_lowercase					(			self:						str    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     MecabTokenizer(mecab_dic="ipadic"    )

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


  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     try:
        _lowerCamelCase   :				str							=     MecabTokenizer(mecab_dic="unidic_lite"    )
     except ModuleNotFoundError:
        return

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


  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     try:
        _lowerCamelCase   :				Any							=     MecabTokenizer(mecab_dic="unidic"    )
     except ModuleNotFoundError:
        return

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


  def 							_lowercase					(			self:						List[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Any							=     MecabTokenizer(do_lower_case=A_							,mecab_dic="ipadic"    )

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


  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     try:
        _lowerCamelCase   :				Any							=     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 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				int							=     MecabTokenizer(normalize_text=A_							,mecab_dic="ipadic"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     self.tokenizer_class(self.vocab_file							,word_tokenizer_type="sudachi"    )
     self.assertIsNotNone(A_    )

     _lowerCamelCase   :				Dict							=     "こんにちは、世界。\nこんばんは、世界。"
     _lowerCamelCase   :				Optional[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[Any]							=     tokenizer_new.tokenize(A_    )

     self.assertListEqual(A_							,A_    )


  @require_sudachi
  def 							_lowercase					(			self:						List[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     SudachiTokenizer(sudachi_dict_type="core"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     SudachiTokenizer(sudachi_dict_type="core"							,sudachi_split_mode="A"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Dict    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     SudachiTokenizer(sudachi_dict_type="core"							,sudachi_split_mode="B"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     _lowerCamelCase   :				str							=     SudachiTokenizer(sudachi_dict_type="core"							,sudachi_split_mode="C"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Union[str, Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Any							=     SudachiTokenizer(do_lower_case=A_							,sudachi_dict_type="core"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				int							=     SudachiTokenizer(normalize_text=A_							,sudachi_dict_type="core"    )

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


  @require_sudachi
  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     SudachiTokenizer(trim_whitespace=A_							,sudachi_dict_type="core"    )

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


  @require_jumanpp
  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[Any]							=     self.tokenizer_class(self.vocab_file							,word_tokenizer_type="jumanpp"    )
     self.assertIsNotNone(A_    )

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

     _lowerCamelCase   :				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   :				List[str]							=     pickle.load(A_    )

     _lowerCamelCase   :				Optional[int]							=     tokenizer_new.tokenize(A_    )

     self.assertListEqual(A_							,A_    )


  @require_jumanpp
  def 							_lowercase					(			self:						Optional[int]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     JumanppTokenizer()

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


  @require_jumanpp
  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Tuple							=     JumanppTokenizer(do_lower_case=A_    )

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


  @require_jumanpp
  def 							_lowercase					(			self:						str    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[Any]							=     JumanppTokenizer(normalize_text=A_    )

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


  @require_jumanpp
  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     JumanppTokenizer(trim_whitespace=A_    )

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


  @require_jumanpp
  def 							_lowercase					(			self:						Dict    ):







     '''simple docstring'''







     _lowerCamelCase   :				Any							=     JumanppTokenizer()

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


  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     ["[UNK]", "[CLS]", "[SEP]", "こんにちは", "こん", "にちは", "ばんは", "##こん", "##にちは", "##ばんは"]

     _lowerCamelCase   :				List[Any]							=     {}
     for i, token in enumerate(A_    ):
        _lowerCamelCase   :				int							=     i
     _lowerCamelCase   :				Any							=     WordpieceTokenizer(vocab=A_							,unk_token="[UNK]"    )

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

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

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

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


  def 							_lowercase					(			self:						Tuple    ):







     '''simple docstring'''







     _lowerCamelCase   :				str							=     BertJapaneseTokenizer.from_pretrained("nlp-waseda/roberta-base-japanese-with-auto-jumanpp"    )
     _lowerCamelCase   :				int							=     tokenizer.subword_tokenizer

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

     _lowerCamelCase   :				List[str]							=     subword_tokenizer.tokenize("こんばんは こんばん にち は こんにちは"    )
     self.assertListEqual(A_							,["▁こん", "ばん", "は", "▁こん", "ばん", "▁に", "ち", "▁は", "▁こんにちは"]    )






  def 							_lowercase					(			self:						Dict    ):







     '''simple docstring'''







     _lowerCamelCase   :				str							=     self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese"    )

     _lowerCamelCase   :				Union[str, Any]							=     tokenizer.encode("ありがとう。"							,add_special_tokens=A_    )
     _lowerCamelCase   :				List[str]							=     tokenizer.encode("どういたしまして。"							,add_special_tokens=A_    )

     _lowerCamelCase   :				str							=     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 							A_      (       __snake_case     ,	unittest.TestCase ):
  lowerCAmelCase__      			=	BertJapaneseTokenizer
  lowerCAmelCase__      			=	False


  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     super().setUp()

     _lowerCamelCase   :				Optional[int]							=     ["[UNK]", "[CLS]", "[SEP]", "こ", "ん", "に", "ち", "は", "ば", "世", "界", "、", "。"]

     _lowerCamelCase   :				Any							=     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 							_lowercase					(			self:						List[Any]							,**__lowerCAmelCase:						Dict    ):







     '''simple docstring'''







     return BertJapaneseTokenizer.from_pretrained(self.tmpdirname							,subword_tokenizer_type="character"							,**A_    )


  def 							_lowercase					(			self:						Any							,__lowerCAmelCase:						Union[str, Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     "こんにちは、世界。 \nこんばんは、世界。"
     _lowerCamelCase   :				Union[str, Any]							=     "こ ん に ち は 、 世 界 。 こ ん ば ん は 、 世 界 。"
     return input_text, output_text


  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						List[str]    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						Union[str, Any]    ):







     '''simple docstring'''







     pass  # TODO add if relevant


  def 							_lowercase					(			self:						Optional[Any]    ):







     '''simple docstring'''







     _lowerCamelCase   :				List[str]							=     self.tokenizer_class(self.vocab_file							,subword_tokenizer_type="character"    )

     _lowerCamelCase   :				Optional[int]							=     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 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Optional[int]							=     ["[UNK]", "[CLS]", "[SEP]", "こ", "ん", "に", "ち", "は", "ば", "世", "界", "、", "。"]

     _lowerCamelCase   :				int							=     {}
     for i, token in enumerate(A_    ):
        _lowerCamelCase   :				str							=     i
     _lowerCamelCase   :				Dict							=     CharacterTokenizer(vocab=A_							,unk_token="[UNK]"    )

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

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

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






  def 							_lowercase					(			self:						int    ):







     '''simple docstring'''







     _lowerCamelCase   :				Any							=     self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese-char"    )

     _lowerCamelCase   :				Optional[int]							=     tokenizer.encode("ありがとう。"							,add_special_tokens=A_    )
     _lowerCamelCase   :				Tuple							=     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 							A_      (       unittest.TestCase ):


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Dict							=     "cl-tohoku/bert-base-japanese"
     _lowerCamelCase   :				List[str]							=     AutoTokenizer.from_pretrained(A_    )
     self.assertIsInstance(A_							,A_    )




class 							A_      (       unittest.TestCase ):


  def 							_lowercase					(			self:						Any    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, 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   :				Dict							=     "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."    )    )






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        MegatronBertForCausalLM,
        MegatronBertForMaskedLM,
        MegatronBertForMultipleChoice,
        MegatronBertForNextSentencePrediction,
        MegatronBertForPreTraining,
        MegatronBertForQuestionAnswering,
        MegatronBertForSequenceClassification,
        MegatronBertForTokenClassification,
        MegatronBertModel,
    )



class    A__  :





      def __init__(					self			,				A_			,				A_=13			,				A_=7			,				A_=True			,				A_=True			,				A_=True			,				A_=True			,				A_=99			,				A_=64			,				A_=32			,				A_=5			,				A_=4			,				A_=37			,				A_="gelu"			,				A_=0.1			,				A_=0.1			,				A_=512			,				A_=16			,				A_=2			,				A_=0.02			,				A_=3			,				A_=4			,				A_=None			,				):
       '''simple docstring'''




       UpperCamelCase      :   Dict			      =  parent
       UpperCamelCase      :   int			      =  batch_size
       UpperCamelCase      :   Dict			      =  seq_length
       UpperCamelCase      :   Union[str, Any]			      =  is_training
       UpperCamelCase      :   int			      =  use_input_mask
       UpperCamelCase      :   Optional[Any]			      =  use_token_type_ids
       UpperCamelCase      :   Optional[Any]			      =  use_labels
       UpperCamelCase      :   str			      =  vocab_size
       UpperCamelCase      :   Union[str, Any]			      =  hidden_size
       UpperCamelCase      :   Any			      =  embedding_size
       UpperCamelCase      :   List[Any]			      =  num_hidden_layers
       UpperCamelCase      :   Optional[int]			      =  num_attention_heads
       UpperCamelCase      :   Optional[Any]			      =  intermediate_size
       UpperCamelCase      :   Dict			      =  hidden_act
       UpperCamelCase      :   Union[str, Any]			      =  hidden_dropout_prob
       UpperCamelCase      :   Optional[Any]			      =  attention_probs_dropout_prob
       UpperCamelCase      :   List[str]			      =  max_position_embeddings
       UpperCamelCase      :   List[Any]			      =  type_vocab_size
       UpperCamelCase      :   Any			      =  type_sequence_label_size
       UpperCamelCase      :   Optional[int]			      =  initializer_range
       UpperCamelCase      :   str			      =  num_labels
       UpperCamelCase      :   List[str]			      =  num_choices
       UpperCamelCase      :   Tuple			      =  scope





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




       UpperCamelCase      :   int			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size  )

       UpperCamelCase      :   List[str]			      =  None
       if self.use_input_mask:
        UpperCamelCase      :   List[str]			      =  random_attention_mask([self.batch_size, self.seq_length]  )

       UpperCamelCase      :   Tuple			      =  None
       if self.use_token_type_ids:
        UpperCamelCase      :   Dict			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.type_vocab_size  )

       UpperCamelCase      :   Dict			      =  None
       UpperCamelCase      :   Union[str, Any]			      =  None
       UpperCamelCase      :   Dict			      =  None
       if self.use_labels:
        UpperCamelCase      :   Optional[Any]			      =  ids_tensor([self.batch_size]			,				self.type_sequence_label_size  )
        UpperCamelCase      :   Tuple			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.num_labels  )
        UpperCamelCase      :   Optional[int]			      =  ids_tensor([self.batch_size]			,				self.num_choices  )

       UpperCamelCase      :   Optional[Any]			      =  self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





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




       return MegatronBertConfig(
           vocab_size=self.vocab_size			,				hidden_size=self.hidden_size			,				num_hidden_layers=self.num_hidden_layers			,				num_attention_heads=self.num_attention_heads			,				intermediate_size=self.intermediate_size			,				embedding_size=self.embedding_size			,				hidden_act=self.hidden_act			,				hidden_dropout_prob=self.hidden_dropout_prob			,				attention_probs_dropout_prob=self.attention_probs_dropout_prob			,				max_position_embeddings=self.max_position_embeddings			,				type_vocab_size=self.type_vocab_size			,				is_decoder=A_			,				initializer_range=self.initializer_range			,				)





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  MegatronBertModel(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Optional[Any]			      =  model(A_			,				attention_mask=A_			,				token_type_ids=A_  )
       UpperCamelCase      :   Dict			      =  model(A_			,				token_type_ids=A_  )
       UpperCamelCase      :   Union[str, Any]			      =  model(A_  )

       self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size)  )
       self.parent.assertEqual(result.pooler_output.shape			,				(self.batch_size, self.hidden_size)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  MegatronBertForMaskedLM(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Optional[int]			      =  model(A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_  )
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  MegatronBertForCausalLM(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Dict			      =  model(A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_  )
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  MegatronBertForNextSentencePrediction(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Dict			      =  model(
           A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_			,				)
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, 2)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  MegatronBertForPreTraining(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   str			      =  model(
           A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_			,				next_sentence_label=A_			,				)
       self.parent.assertEqual(result.prediction_logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size)  )
       self.parent.assertEqual(result.seq_relationship_logits.shape			,				(self.batch_size, 2)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   str			      =  MegatronBertForQuestionAnswering(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Union[str, Any]			      =  model(
           A_			,				attention_mask=A_			,				token_type_ids=A_			,				start_positions=A_			,				end_positions=A_			,				)
       self.parent.assertEqual(result.start_logits.shape			,				(self.batch_size, self.seq_length)  )
       self.parent.assertEqual(result.end_logits.shape			,				(self.batch_size, self.seq_length)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  self.num_labels
       UpperCamelCase      :   Optional[int]			      =  MegatronBertForSequenceClassification(A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   List[Any]			      =  model(A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_  )
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.num_labels)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  self.num_labels
       UpperCamelCase      :   List[str]			      =  MegatronBertForTokenClassification(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   List[str]			      =  model(A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_  )
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.num_labels)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  self.num_choices
       UpperCamelCase      :   int			      =  MegatronBertForMultipleChoice(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Optional[int]			      =  input_ids.unsqueeze(1  ).expand(-1			,				self.num_choices			,				-1  ).contiguous()
       UpperCamelCase      :   Tuple			      =  token_type_ids.unsqueeze(1  ).expand(-1			,				self.num_choices			,				-1  ).contiguous()
       UpperCamelCase      :   List[str]			      =  input_mask.unsqueeze(1  ).expand(-1			,				self.num_choices			,				-1  ).contiguous()
       UpperCamelCase      :   Tuple			      =  model(
           A_			,				attention_mask=A_			,				token_type_ids=A_			,				labels=A_			,				)
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.num_choices)  )





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




       UpperCamelCase      :   Optional[Any]			      =  self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       (
           UpperCamelCase
       )      ,       
       )      :   Any			      =  config_and_inputs
       UpperCamelCase      :   Optional[int]			      =  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_torch
class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :Tuple		       =  (
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _UpperCAmelCase      :Optional[Any]		       =  (
          {
              'feature-extraction': MegatronBertModel,
              'fill-mask': MegatronBertForMaskedLM,
              'question-answering': MegatronBertForQuestionAnswering,
              'text-classification': MegatronBertForSequenceClassification,
              'text-generation': MegatronBertForCausalLM,
              'token-classification': MegatronBertForTokenClassification,
              'zero-shot': MegatronBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _UpperCAmelCase      :Optional[Any]		       =  True
      # test_resize_embeddings = False
      _UpperCAmelCase      :Optional[Any]		       =  False





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_=False  ):
       '''simple docstring'''




       UpperCamelCase      :   Any			      =  super()._prepare_for_class(A_			,				A_			,				return_labels=A_  )

       if return_labels:
        if model_class in get_values(A_  ):
         UpperCamelCase      :   str			      =  torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length)			,				dtype=torch.long			,				device=A_  )
         UpperCamelCase      :   str			      =  torch.zeros(
             self.model_tester.batch_size			,				dtype=torch.long			,				device=A_  )
       return inputs_dict





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




       UpperCamelCase      :   Optional[int]			      =  MegatronBertModelTester(self  )
       UpperCamelCase      :   Optional[int]			      =  ConfigTester(self			,				config_class=A_			,				hidden_size=37  )





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




       self.config_tester.run_common_tests()





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




       UpperCamelCase      :   List[str]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*A_  )





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




       UpperCamelCase      :   int			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A_  )





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




       UpperCamelCase      :   Optional[Any]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A_  )





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




       UpperCamelCase      :   Any			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A_  )





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




       UpperCamelCase      :   int			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*A_  )





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




       UpperCamelCase      :   Dict			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*A_  )





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




       UpperCamelCase      :   Union[str, Any]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A_  )





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




       UpperCamelCase      :   Optional[int]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*A_  )



def 						A_							(	_lowerCAmelCase	)		->					int:
 return torch.tensor(
     _lowerCAmelCase				,  dtype=torch.long				,  device=_lowerCAmelCase				,  )


__lowerCamelCase			:	Optional[int]					    =	1E-4



@require_torch
@require_sentencepiece
@require_tokenizers
class    A__  (			unittest.TestCase     ):





      @slow
      @unittest.skip("Model is not available."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  "nvidia/megatron-bert-uncased-345m"
       if "MYDIR" in os.environ:
        UpperCamelCase      :   Union[str, Any]			      =  os.path.join(os.environ["MYDIR"]			,				A_  )
       UpperCamelCase      :   List[Any]			      =  MegatronBertModel.from_pretrained(A_  )
       model.to(A_  )
       model.half()
       UpperCamelCase      :   Optional[Any]			      =  _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]  )
       with torch.no_grad():
        UpperCamelCase      :   Tuple			      =  model(A_  )[0]
       UpperCamelCase      :   Dict			      =  torch.Size((1, 9, 1024)  )
       self.assertEqual(output.shape			,				A_  )

       UpperCamelCase      :   Union[str, Any]			      =  [-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
       for ii in range(3  ):
        for jj in range(3  ):
         UpperCamelCase      :   List[str]			      =  output[0, ii, jj]
         UpperCamelCase      :   List[Any]			      =  expected[3 * ii + jj]
         UpperCamelCase      :   Optional[Any]			      =  "ii={} jj={} a={} b={}".format(A_			,				A_			,				A_			,				A_  )
         self.assertTrue(math.isclose(A_			,				A_			,				rel_tol=A_			,				abs_tol=A_  )			,				msg=A_  )
