

'''simple docstring'''

import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
    import os

    # The slow tests are often failing with OOM error on GPU
    # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
    # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    A_      =							"platform"

    import jax
    import jax.numpy as jnp

    from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer





def 				A_					(	snake_case ,    snake_case ,    snake_case=None ,    snake_case=None	):
  if attention_mask is None:
    SCREAMING_SNAKE_CASE:Tuple   =   np.where(input_ids != config.pad_token_id ,    1 ,    0	)
  return {
      "input_ids": input_ids,
      "attention_mask": attention_mask,
  }





@require_flax
class 							_snake_case	:
   def __init__(	self     : List[Any]       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : List[str]=13       ,SCREAMING_SNAKE_CASE__     : List[str]=7       ,SCREAMING_SNAKE_CASE__     : Optional[int]=True       ,SCREAMING_SNAKE_CASE__     : List[Any]=False       ,SCREAMING_SNAKE_CASE__     : Optional[int]=99       ,SCREAMING_SNAKE_CASE__     : Optional[int]=16       ,SCREAMING_SNAKE_CASE__     : Any=2       ,SCREAMING_SNAKE_CASE__     : Tuple=4       ,SCREAMING_SNAKE_CASE__     : List[str]=4       ,SCREAMING_SNAKE_CASE__     : int="gelu"       ,SCREAMING_SNAKE_CASE__     : Optional[int]=0.1       ,SCREAMING_SNAKE_CASE__     : str=0.1       ,SCREAMING_SNAKE_CASE__     : str=20       ,SCREAMING_SNAKE_CASE__     : List[str]=2       ,SCREAMING_SNAKE_CASE__     : Dict=1       ,SCREAMING_SNAKE_CASE__     : str=0       ,SCREAMING_SNAKE_CASE__     : List[Any]=16       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]=16       ,SCREAMING_SNAKE_CASE__     : Union[str, Any]=0.02       ,):
     SCREAMING_SNAKE_CASE:Optional[Any]   =   parent
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   batch_size
     SCREAMING_SNAKE_CASE:str   =   seq_length
     SCREAMING_SNAKE_CASE:Dict   =   is_training
     SCREAMING_SNAKE_CASE:str   =   use_labels
     SCREAMING_SNAKE_CASE:List[Any]   =   vocab_size
     SCREAMING_SNAKE_CASE:Tuple   =   hidden_size
     SCREAMING_SNAKE_CASE:Any   =   num_hidden_layers
     SCREAMING_SNAKE_CASE:str   =   num_attention_heads
     SCREAMING_SNAKE_CASE:Any   =   intermediate_size
     SCREAMING_SNAKE_CASE:Optional[int]   =   hidden_act
     SCREAMING_SNAKE_CASE:int   =   hidden_dropout_prob
     SCREAMING_SNAKE_CASE:List[str]   =   attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE:str   =   max_position_embeddings
     SCREAMING_SNAKE_CASE:Optional[int]   =   eos_token_id
     SCREAMING_SNAKE_CASE:Tuple   =   pad_token_id
     SCREAMING_SNAKE_CASE:List[str]   =   bos_token_id
     SCREAMING_SNAKE_CASE:Optional[int]   =   embed_dim
     SCREAMING_SNAKE_CASE:str   =   word_embed_proj_dim
     SCREAMING_SNAKE_CASE:Optional[int]   =   initializer_range
     SCREAMING_SNAKE_CASE:List[str]   =   False
   def       __UpperCamelCase  (	self     : str							):
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       ,self.vocab_size							)       ,3       ,self.vocab_size							)
     SCREAMING_SNAKE_CASE:List[Any]   =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       ,dtype=np.intaa							))       ,-1							)

     SCREAMING_SNAKE_CASE:List[str]   =   OPTConfig(
         vocab_size=self.vocab_size       ,hidden_size=self.hidden_size       ,num_hidden_layers=self.num_hidden_layers       ,num_attention_heads=self.num_attention_heads       ,ffn_dim=self.intermediate_size       ,dropout=self.hidden_dropout_prob       ,attention_dropout=self.attention_probs_dropout_prob       ,max_position_embeddings=self.max_position_embeddings       ,eos_token_id=self.eos_token_id       ,bos_token_id=self.bos_token_id       ,pad_token_id=self.pad_token_id       ,embed_dim=self.embed_dim       ,is_encoder_decoder=SCREAMING_SNAKE_CASE__       ,word_embed_proj_dim=self.word_embed_proj_dim       ,initializer_range=self.initializer_range       ,use_cache=SCREAMING_SNAKE_CASE__       ,)
     SCREAMING_SNAKE_CASE:int   =   prepare_opt_inputs_dict(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
     return config, inputs_dict
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:Dict   =   self.prepare_config_and_inputs()
     return config, inputs_dict
   def       __UpperCamelCase  (	self     : Any       ,SCREAMING_SNAKE_CASE__     : Tuple       ,SCREAMING_SNAKE_CASE__     : List[str]       ,SCREAMING_SNAKE_CASE__     : List[Any]							):
     SCREAMING_SNAKE_CASE:Dict   =   20
     SCREAMING_SNAKE_CASE:Optional[int]   =   model_class_name(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:List[Any]   =   inputs_dict["input_ids"]
     SCREAMING_SNAKE_CASE:str   =   inputs_dict["attention_mask"]

     SCREAMING_SNAKE_CASE:Optional[Any]   =   model.init_cache(input_ids.shape[0]       ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   jnp.ones((input_ids.shape[0], max_length)       ,dtype="i4"							)

     SCREAMING_SNAKE_CASE:Optional[int]   =   jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1							)[None, :]       ,(input_ids.shape[0], input_ids.shape[-1] - 1)       ,)
     SCREAMING_SNAKE_CASE:str   =   model(
         input_ids[:, :-1]       ,attention_mask=SCREAMING_SNAKE_CASE__       ,past_key_values=SCREAMING_SNAKE_CASE__       ,position_ids=SCREAMING_SNAKE_CASE__       ,)

     SCREAMING_SNAKE_CASE:Tuple   =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,dtype="i4"							)
     SCREAMING_SNAKE_CASE:Dict   =   model(
         input_ids[:, -1:]       ,attention_mask=SCREAMING_SNAKE_CASE__       ,past_key_values=outputs_cache.past_key_values       ,position_ids=SCREAMING_SNAKE_CASE__       ,)

     SCREAMING_SNAKE_CASE:Dict   =   model(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
     self.parent.assertTrue(diff < 1e-3       ,msg=F'''Max diff is {diff}'''							)







   def       __UpperCamelCase  (	self     : Any       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : Optional[Any]       ,SCREAMING_SNAKE_CASE__     : List[Any]							):
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   20
     SCREAMING_SNAKE_CASE:Any   =   model_class_name(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:List[str]   =   (
         inputs_dict["input_ids"],
         inputs_dict["attention_mask"],
     )

     SCREAMING_SNAKE_CASE:List[str]   =   jnp.concatenate(
         [
             attention_mask,
             jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])							),
         ]       ,axis=-1       ,)

     SCREAMING_SNAKE_CASE:Tuple   =   model.init_cache(input_ids.shape[0]       ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Tuple   =   jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1							)[None, :]       ,(input_ids.shape[0], input_ids.shape[-1] - 1)       ,)

     SCREAMING_SNAKE_CASE:Dict   =   model(
         input_ids[:, :-1]       ,attention_mask=SCREAMING_SNAKE_CASE__       ,past_key_values=SCREAMING_SNAKE_CASE__       ,position_ids=SCREAMING_SNAKE_CASE__       ,)
     SCREAMING_SNAKE_CASE:Tuple   =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,dtype="i4"							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   model(
         input_ids[:, -1:]       ,past_key_values=outputs_cache.past_key_values       ,attention_mask=SCREAMING_SNAKE_CASE__       ,position_ids=SCREAMING_SNAKE_CASE__       ,)

     SCREAMING_SNAKE_CASE:Optional[int]   =   model(SCREAMING_SNAKE_CASE__       ,attention_mask=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Tuple   =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
     self.parent.assertTrue(diff < 1e-3       ,msg=F'''Max diff is {diff}'''							)





@require_flax
class 							_snake_case	(    _a		,					unittest.TestCase		,					_a  ):
   _A    :       Any      				=   (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
   _A    :       List[str]      				=   () if is_flax_available() else ()
   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE:List[str]   =   FlaxOPTModelTester(self							)
   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:List[Any]   =   self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
       self.model_tester.check_use_cache_forward(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Optional[Any]							):
     SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:Optional[Any]   =   self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
       self.model_tester.check_use_cache_forward_with_attn_mask(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)







   @slow
   def       __UpperCamelCase  (	self     : Optional[Any]							):
     for model_class_name in self.all_model_classes:
       SCREAMING_SNAKE_CASE:List[Any]   =   model_class_name.from_pretrained("facebook/opt-125m"							)
       SCREAMING_SNAKE_CASE:Tuple   =   np.ones((1, 1)							) * model.config.eos_token_id
       SCREAMING_SNAKE_CASE:List[Any]   =   model(SCREAMING_SNAKE_CASE__							)
       self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)





@require_sentencepiece
@require_flax
class 							_snake_case	(    unittest.TestCase  ):
   @slow
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:int   =   FlaxOPTModel.from_pretrained("facebook/opt-350m"							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   jnp.array([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]							)
     SCREAMING_SNAKE_CASE:Any   =   model(input_ids=SCREAMING_SNAKE_CASE__							).last_hidden_state
     SCREAMING_SNAKE_CASE:Optional[int]   =   (1, 11, 512)
     self.assertEqual(output.shape       ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:str   =   jnp.array(
         [[-0.2_867, -1.9_256, -0.3_062], [-1.2_711, -0.1_337, -0.1_897], [0.4_109, 0.1_187, -1.3_142]]							)
     self.assertTrue(jnp.allclose(output[:, :3, :3]       ,SCREAMING_SNAKE_CASE__       ,atol=4e-2							)							)





@require_flax
@slow
class 							_snake_case	(    unittest.TestCase  ):
   def       __UpperCamelCase  (	self     : Dict							):
     super().setUp()
     SCREAMING_SNAKE_CASE:Any   =   "facebook/opt-350m"







   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE:Dict   =   FlaxOPTForCausalLM.from_pretrained(self.path_model							)
     SCREAMING_SNAKE_CASE:Tuple   =   GPTaTokenizer.from_pretrained(self.path_model							)

     SCREAMING_SNAKE_CASE:Any   =   [
         "Today is a beautiful day and I want to",
         "In the city of",
         "Paris is the capital of France and",
         "Computers and mobile phones have taken",
     ]
     # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
     SCREAMING_SNAKE_CASE:List[str]   =   tokenizer(SCREAMING_SNAKE_CASE__       ,return_tensors="jax"       ,padding=SCREAMING_SNAKE_CASE__       ,add_special_tokens=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:List[str]   =   model(inputs.input_ids       ,attention_mask=inputs.attention_mask							)[0].mean(axis=-1							)
     SCREAMING_SNAKE_CASE:str   =   jnp.array(
         [
             [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
             [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
             [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
             [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
         ]							)
     self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,atol=4e-2							)							)

     SCREAMING_SNAKE_CASE:List[Any]   =   jax.jit(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   model(inputs.input_ids       ,attention_mask=inputs.attention_mask							)[0].mean(axis=-1							)
     self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,atol=4e-2							)							)





@require_flax
@slow
class 							_snake_case	(    unittest.TestCase  ):
   @property
   def       __UpperCamelCase  (	self     : List[Any]							):
     return [
         "Today is a beautiful day and I want",
         "In the city of",
         "Paris is the capital of France and",
         "Computers and mobile phones have taken",
     ]
   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE:str   =   "facebook/opt-125m"

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   [
         "Today is a beautiful day and I want to",
         "In the city of New York, the city",
         "Paris is the capital of France and the capital",
         "Computers and mobile phones have taken over the",
     ]

     SCREAMING_SNAKE_CASE:int   =   []

     SCREAMING_SNAKE_CASE:int   =   FlaxOPTForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict   =   GPTaTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__							)

     for prompt in self.prompts:
       SCREAMING_SNAKE_CASE:Tuple   =   tokenizer(SCREAMING_SNAKE_CASE__       ,return_tensors="jax"							).input_ids

       SCREAMING_SNAKE_CASE:List[Any]   =   model.generate(SCREAMING_SNAKE_CASE__       ,max_length=10							)
       SCREAMING_SNAKE_CASE:Union[str, Any]   =   generated_ids[0]

       SCREAMING_SNAKE_CASE:List[str]   =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE__       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)
       predicted_outputs += generated_string

     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : List[str]							):
     SCREAMING_SNAKE_CASE:Optional[int]   =   "facebook/opt-350m"

     SCREAMING_SNAKE_CASE:List[Any]   =   [
         "Today is a beautiful day and I want to",
         "In the city of San Francisco, the city",
         "Paris is the capital of France and the capital",
         "Computers and mobile phones have taken over the",
     ]

     SCREAMING_SNAKE_CASE:List[str]   =   []
     SCREAMING_SNAKE_CASE:Dict   =   FlaxOPTForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:List[str]   =   GPTaTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__							)

     for prompt in self.prompts:
       SCREAMING_SNAKE_CASE:Optional[Any]   =   tokenizer(SCREAMING_SNAKE_CASE__       ,return_tensors="jax"							).input_ids

       SCREAMING_SNAKE_CASE:Optional[int]   =   model.generate(SCREAMING_SNAKE_CASE__       ,max_length=10							)
       SCREAMING_SNAKE_CASE:List[str]   =   generated_ids[0]

       SCREAMING_SNAKE_CASE:Union[str, Any]   =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE__       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)
       predicted_outputs += generated_string

     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Optional[int]							):
     SCREAMING_SNAKE_CASE:Optional[int]   =   "facebook/opt-125m"
     SCREAMING_SNAKE_CASE:Any   =   [
         "Today is a beautiful day and I want to thank",
         "In the city of Rome Canaver Canaver Canaver Canaver",
     ]
     SCREAMING_SNAKE_CASE:Optional[int]   =   FlaxOPTForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Any   =   GPTaTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   tokenizer(
         [
             "Today is a beautiful day and I want to",
             "In the city of",
         ]       ,return_tensors="jax"       ,padding=SCREAMING_SNAKE_CASE__       ,)

     SCREAMING_SNAKE_CASE:Dict   =   jax.jit(model.generate							)

     SCREAMING_SNAKE_CASE:Any   =   jit_generate(inputs["input_ids"]       ,attention_mask=inputs["attention_mask"]							).sequences

     SCREAMING_SNAKE_CASE:List[Any]   =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE__       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)

     self.assertIsNotNone(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)







   def       __UpperCamelCase  (	self     : Tuple							):
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   "facebook/opt-350m"

     SCREAMING_SNAKE_CASE:Optional[Any]   =   GPTaTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Tuple   =   FlaxOPTForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Tuple   =   "left"

     # use different length sentences to test batching
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   [
         "Hello, my dog is a little",
         "Today, I",
     ]

     SCREAMING_SNAKE_CASE:Optional[Any]   =   tokenizer(SCREAMING_SNAKE_CASE__       ,return_tensors="jax"       ,padding=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict   =   inputs["input_ids"]

     SCREAMING_SNAKE_CASE:List[Any]   =   model.generate(input_ids=SCREAMING_SNAKE_CASE__       ,attention_mask=inputs["attention_mask"]       ,trace=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   tokenizer(sentences[0]       ,return_tensors="jax"							).input_ids
     SCREAMING_SNAKE_CASE:str   =   model.generate(input_ids=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Any   =   inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
     SCREAMING_SNAKE_CASE:Optional[int]   =   tokenizer(sentences[1]       ,return_tensors="jax"							).input_ids
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   model.generate(input_ids=SCREAMING_SNAKE_CASE__       ,max_length=model.config.max_length - num_paddings							)

     SCREAMING_SNAKE_CASE:Any   =   tokenizer.batch_decode(outputs[0]       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int   =   tokenizer.decode(output_non_padded[0][0]       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict   =   tokenizer.decode(output_padded[0][0]       ,skip_special_tokens=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Optional[int]   =   [
         "Hello, my dog is a little bit of a dork.\nI'm a little bit",
         "Today, I was in the middle of a conversation with a friend about the",
     ]
     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,[non_padded_sentence, padded_sentence]							)









'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


A_      =							{"configuration_yoso": ["YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP", "YosoConfig"]}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A_      =							[
        "YOSO_PRETRAINED_MODEL_ARCHIVE_LIST",
        "YosoForMaskedLM",
        "YosoForMultipleChoice",
        "YosoForQuestionAnswering",
        "YosoForSequenceClassification",
        "YosoForTokenClassification",
        "YosoLayer",
        "YosoModel",
        "YosoPreTrainedModel",
    ]


if TYPE_CHECKING:
    from .configuration_yoso import YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP, YosoConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_yoso import (
            YOSO_PRETRAINED_MODEL_ARCHIVE_LIST,
            YosoForMaskedLM,
            YosoForMultipleChoice,
            YosoForQuestionAnswering,
            YosoForSequenceClassification,
            YosoForTokenClassification,
            YosoLayer,
            YosoModel,
            YosoPreTrainedModel,
        )


else:
    import sys

    A_      =							_LazyModule(__name__, globals()["__file__"], _import_structure)






