





'''simple docstring'''



import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY


@is_pipeline_test
@require_torch_or_tf
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Optional[int]      =   MODEL_FOR_CAUSAL_LM_MAPPING
     UpperCAmelCase__       :						Tuple      =   TF_MODEL_FOR_CAUSAL_LM_MAPPING

     @require_torch
     def 	UpperCamelCase(					self	):
       _snake_case     		=				pipeline(task="text-generation"						,   model="sshleifer/tiny-ctrl"						,   framework="pt"	)
       # Using `do_sample=False` to force deterministic output
       _snake_case     		=				text_generator("This is a test"						,   do_sample=lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [
               {
                   "generated_text": (
                       "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                       " oscope. FiliFili@@"
                   )
               }
           ]						,   )

       _snake_case     		=				text_generator(["This is a test", "This is a second test"]	)
       self.assertEqual(
           lowerCamelCase						,   [
               [
                   {
                       "generated_text": (
                           "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                           " oscope. FiliFili@@"
                       )
                   }
               ],
               [
                   {
                       "generated_text": (
                           "This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy"
                           " oscope. oscope. FiliFili@@"
                       )
                   }
               ],
           ]						,   )

       _snake_case     		=				text_generator("This is a test"						,   do_sample=lowerCamelCase						,   num_return_sequences=2						,   return_tensors=lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [
               {"generated_token_ids": ANY(lowerCamelCase	)},
               {"generated_token_ids": ANY(lowerCamelCase	)},
           ]						,   )
       _snake_case     		=				text_generator.model.config.eos_token_id
       _snake_case     		=				"<pad>"
       _snake_case     		=				text_generator(
           ["This is a test", "This is a second test"]						,   do_sample=lowerCamelCase						,   num_return_sequences=2						,   batch_size=2						,   return_tensors=lowerCamelCase						,   )
       self.assertEqual(
           lowerCamelCase						,   [
               [
                   {"generated_token_ids": ANY(lowerCamelCase	)},
                   {"generated_token_ids": ANY(lowerCamelCase	)},
               ],
               [
                   {"generated_token_ids": ANY(lowerCamelCase	)},
                   {"generated_token_ids": ANY(lowerCamelCase	)},
               ],
           ]						,   )

     @require_tf
     def 	UpperCamelCase(					self	):
       _snake_case     		=				pipeline(task="text-generation"						,   model="sshleifer/tiny-ctrl"						,   framework="tf"	)

       # Using `do_sample=False` to force deterministic output
       _snake_case     		=				text_generator("This is a test"						,   do_sample=lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [
               {
                   "generated_text": (
                       "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                       " please,"
                   )
               }
           ]						,   )

       _snake_case     		=				text_generator(["This is a test", "This is a second test"]						,   do_sample=lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [
               [
                   {
                       "generated_text": (
                           "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                           " please,"
                       )
                   }
               ],
               [
                   {
                       "generated_text": (
                           "This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes"
                           " Cannes 閲閲Cannes Cannes Cannes 攵 please,"
                       )
                   }
               ],
           ]						,   )

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				TextGenerationPipeline(model=lowerCamelCase						,   tokenizer=lowerCamelCase	)
       return text_generator, ["This is a test", "Another test"]

     def 	UpperCamelCase(					self	):
       _snake_case     		=				"Hello I believe in"
       _snake_case     		=				pipeline("text-generation"						,   model="hf-internal-testing/tiny-random-gpt2"	)
       _snake_case     		=				text_generator(lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}]						,   )

       _snake_case     		=				text_generator(lowerCamelCase						,   stop_sequence=" fe"	)
       self.assertEqual(lowerCamelCase						,   [{"generated_text": "Hello I believe in fe"}]	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				text_generator.model
       _snake_case     		=				text_generator.tokenizer

       _snake_case     		=				text_generator("This is a test"	)
       self.assertEqual(lowerCamelCase						,   [{"generated_text": ANY(lowerCamelCase	)}]	)
       self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"	)	)

       _snake_case     		=				text_generator("This is a test"						,   return_full_text=lowerCamelCase	)
       self.assertEqual(lowerCamelCase						,   [{"generated_text": ANY(lowerCamelCase	)}]	)
       self.assertNotIn("This is a test"						,   outputs[0]["generated_text"]	)

       _snake_case     		=				pipeline(task="text-generation"						,   model=lowerCamelCase						,   tokenizer=lowerCamelCase						,   return_full_text=lowerCamelCase	)
       _snake_case     		=				text_generator("This is a test"	)
       self.assertEqual(lowerCamelCase						,   [{"generated_text": ANY(lowerCamelCase	)}]	)
       self.assertNotIn("This is a test"						,   outputs[0]["generated_text"]	)

       _snake_case     		=				text_generator("This is a test"						,   return_full_text=lowerCamelCase	)
       self.assertEqual(lowerCamelCase						,   [{"generated_text": ANY(lowerCamelCase	)}]	)
       self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"	)	)

       _snake_case     		=				text_generator(["This is great !", "Something else"]						,   num_return_sequences=2						,   do_sample=lowerCamelCase	)
       self.assertEqual(
           lowerCamelCase						,   [
               [{"generated_text": ANY(lowerCamelCase	)}, {"generated_text": ANY(lowerCamelCase	)}],
               [{"generated_text": ANY(lowerCamelCase	)}, {"generated_text": ANY(lowerCamelCase	)}],
           ]						,   )

       if text_generator.tokenizer.pad_token is not None:
         _snake_case     		=				text_generator(
             ["This is great !", "Something else"]						,   num_return_sequences=2						,   batch_size=2						,   do_sample=lowerCamelCase	)
         self.assertEqual(
             lowerCamelCase						,   [
                 [{"generated_text": ANY(lowerCamelCase	)}, {"generated_text": ANY(lowerCamelCase	)}],
                 [{"generated_text": ANY(lowerCamelCase	)}, {"generated_text": ANY(lowerCamelCase	)}],
             ]						,   )

       with self.assertRaises(lowerCamelCase	):
         _snake_case     		=				text_generator("test"						,   return_full_text=lowerCamelCase						,   return_text=lowerCamelCase	)
       with self.assertRaises(lowerCamelCase	):
         _snake_case     		=				text_generator("test"						,   return_full_text=lowerCamelCase						,   return_tensors=lowerCamelCase	)
       with self.assertRaises(lowerCamelCase	):
         _snake_case     		=				text_generator("test"						,   return_text=lowerCamelCase						,   return_tensors=lowerCamelCase	)

       # Empty prompt is slighly special
       # it requires BOS token to exist.
       # Special case for Pegasus which will always append EOS so will
       # work even without BOS.
       if (
           text_generator.tokenizer.bos_token_id is not None
           or "Pegasus" in tokenizer.__class__.__name__
           or "Git" in model.__class__.__name__
       ):
         _snake_case     		=				text_generator(""	)
         self.assertEqual(lowerCamelCase						,   [{"generated_text": ANY(lowerCamelCase	)}]	)
       else:
         with self.assertRaises((ValueError, AssertionError)	):
           _snake_case     		=				text_generator(""	)

       if text_generator.framework == "tf":
         # TF generation does not support max_new_tokens, and it's impossible
         # to control long generation with only max_length without
         # fancy calculation, dismissing tests for now.
         return
       # We don't care about infinite range models.
       # They already work.
       # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
       _snake_case     		=				["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
       if (
           tokenizer.model_max_length < 10_000
           and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
       ):
         # Handling of large generations
         with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)	):
           text_generator("This is a test" * 500						,   max_new_tokens=20	)

         _snake_case     		=				text_generator("This is a test" * 500						,   handle_long_generation="hole"						,   max_new_tokens=20	)
         # Hole strategy cannot work
         with self.assertRaises(lowerCamelCase	):
           text_generator(
               "This is a test" * 500						,   handle_long_generation="hole"						,   max_new_tokens=tokenizer.model_max_length + 10						,   )

     @require_torch
     @require_accelerate
     @require_torch_gpu
     def 	UpperCamelCase(					self	):
       import torch

       # Classic `model_kwargs`
       _snake_case     		=				pipeline(
           model="hf-internal-testing/tiny-random-bloom"						,   model_kwargs={"device_map": "auto", "torch_dtype": torch.bfloataa}						,   )
       self.assertEqual(pipe.model.device						,   torch.device(0	)	)
       self.assertEqual(pipe.model.lm_head.weight.dtype						,   torch.bfloataa	)
       _snake_case     		=				pipe("This is a test"	)
       self.assertEqual(
           lowerCamelCase						,   [
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]						,   )

       # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
       _snake_case     		=				pipeline(model="hf-internal-testing/tiny-random-bloom"						,   device_map="auto"						,   torch_dtype=torch.bfloataa	)
       self.assertEqual(pipe.model.device						,   torch.device(0	)	)
       self.assertEqual(pipe.model.lm_head.weight.dtype						,   torch.bfloataa	)
       _snake_case     		=				pipe("This is a test"	)
       self.assertEqual(
           lowerCamelCase						,   [
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]						,   )

       # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
       _snake_case     		=				pipeline(model="hf-internal-testing/tiny-random-bloom"						,   device_map="auto"	)
       self.assertEqual(pipe.model.device						,   torch.device(0	)	)
       self.assertEqual(pipe.model.lm_head.weight.dtype						,   torch.floataa	)
       _snake_case     		=				pipe("This is a test"	)
       self.assertEqual(
           lowerCamelCase						,   [
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]						,   )

     @require_torch
     @require_torch_gpu
     def 	UpperCamelCase(					self	):
       import torch

       _snake_case     		=				pipeline(model="hf-internal-testing/tiny-random-bloom"						,   device=0						,   torch_dtype=torch.floataa	)
       pipe("This is a test"	)

     @require_torch
     @require_accelerate
     @require_torch_gpu
     def 	UpperCamelCase(					self	):
       import torch

       _snake_case     		=				pipeline(model="hf-internal-testing/tiny-random-bloom"						,   device_map="auto"						,   torch_dtype=torch.floataa	)
       pipe("This is a test"						,   do_sample=lowerCamelCase						,   top_p=0.5	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				"Hello world"
       _snake_case     		=				pipeline("text-generation"						,   model="hf-internal-testing/tiny-random-gpt2"	)
       if text_generator.model.framework == "tf":
         _snake_case     		=				logging.get_logger("transformers.generation.tf_utils"	)
       else:
         _snake_case     		=				logging.get_logger("transformers.generation.utils"	)
       _snake_case     		=				"Both `max_new_tokens`"  # The beggining of the message to be checked in this test

       # Both are set by the user -> log warning
       with CaptureLogger(lowerCamelCase	) as cl:
         _snake_case     		=				text_generator(lowerCamelCase						,   max_length=10						,   max_new_tokens=1	)
       self.assertIn(lowerCamelCase						,   cl.out	)

       # The user only sets one -> no warning
       with CaptureLogger(lowerCamelCase	) as cl:
         _snake_case     		=				text_generator(lowerCamelCase						,   max_new_tokens=1	)
       self.assertNotIn(lowerCamelCase						,   cl.out	)

       with CaptureLogger(lowerCamelCase	) as cl:
         _snake_case     		=				text_generator(lowerCamelCase						,   max_length=10	)
       self.assertNotIn(lowerCamelCase						,   cl.out	)









'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__magic_name__   :						Any    =			logging.get_logger(__name__)

__magic_name__   :						Dict    =			{
    """google/bit-50""": """https://huggingface.co/google/bit-50/resolve/main/config.json""",
}


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						List[str]      =   '''bit'''
     UpperCAmelCase__       :						Union[str, Any]      =   ['''preactivation''', '''bottleneck''']
     UpperCAmelCase__       :						int      =   ['''SAME''', '''VALID''']

     def __init__(					self						,   lowerCamelCase=3						,   lowerCamelCase=64						,   lowerCamelCase=[256, 512, 1_024, 2_048]						,   lowerCamelCase=[3, 4, 6, 3]						,   lowerCamelCase="preactivation"						,   lowerCamelCase="relu"						,   lowerCamelCase=None						,   lowerCamelCase=32						,   lowerCamelCase=0.0						,   lowerCamelCase=False						,   lowerCamelCase=32						,   lowerCamelCase=1						,   lowerCamelCase=None						,   lowerCamelCase=None						,   **lowerCamelCase						,   ):
       super().__init__(**lowerCamelCase	)
       if layer_type not in self.layer_types:
         raise ValueError(F'''layer_type={layer_type} is not one of {",".join(self.layer_types	)}'''	)
       if global_padding is not None:
         if global_padding.upper() in self.supported_padding:
           _snake_case     		=				global_padding.upper()
         else:
           raise ValueError(F'''Padding strategy {global_padding} not supported'''	)
       _snake_case     		=				num_channels
       _snake_case     		=				embedding_size
       _snake_case     		=				hidden_sizes
       _snake_case     		=				depths
       _snake_case     		=				layer_type
       _snake_case     		=				hidden_act
       _snake_case     		=				global_padding
       _snake_case     		=				num_groups
       _snake_case     		=				drop_path_rate
       _snake_case     		=				embedding_dynamic_padding
       _snake_case     		=				output_stride
       _snake_case     		=				width_factor

       _snake_case     		=				["stem"] + [F'''stage{idx}''' for idx in range(1						,   len(lowerCamelCase	) + 1	)]
       _snake_case		,		_snake_case     		=				get_aligned_output_features_output_indices(
           out_features=lowerCamelCase						,   out_indices=lowerCamelCase						,   stage_names=self.stage_names	)


