
"""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 __lowerCAmelCase	(   unittest.TestCase    ):
    lowercase		     =		MODEL_FOR_CAUSAL_LM_MAPPING
    lowercase		     =		TF_MODEL_FOR_CAUSAL_LM_MAPPING





    @require_torch
    def 		UpperCAmelCase      (			self   ):

     '''simple docstring'''




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

     __UpperCamelCase       						= text_generator(['This is a test', 'This is a second test']   )
     self.assertEqual(
         __a			,  [
             [
                 {
                     '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@@'
                     )
                 }
             ],
         ]			,  )

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





    @require_tf
    def 		UpperCAmelCase      (			self   ):

     '''simple docstring'''




     __UpperCamelCase       						= pipeline(task='text-generation'			,  model='sshleifer/tiny-ctrl'			,  framework='tf'   )

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

     __UpperCamelCase       						= text_generator(['This is a test', 'This is a second test']			,  do_sample=__a   )
     self.assertEqual(
         __a			,  [
             [
                 {
                     '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			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

     '''simple docstring'''




     __UpperCamelCase       						= TextGenerationPipeline(model=__a			,  tokenizer=__a   )
     return text_generator, ["This is a test", "Another test"]





    def 		UpperCAmelCase      (			self   ):

     '''simple docstring'''




     __UpperCamelCase       						= 'Hello I believe in'
     __UpperCamelCase       						= pipeline('text-generation'			,  model='hf-internal-testing/tiny-random-gpt2'   )
     __UpperCamelCase       						= text_generator(__a   )
     self.assertEqual(
         __a			,  [{'generated_text': 'Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'}]			,  )

     __UpperCamelCase       						= text_generator(__a			,  stop_sequence=' fe'   )
     self.assertEqual(__a			,  [{'generated_text': 'Hello I believe in fe'}]   )





    def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

     '''simple docstring'''




     __UpperCamelCase       						= text_generator.model
     __UpperCamelCase       						= text_generator.tokenizer

     __UpperCamelCase       						= text_generator('This is a test'   )
     self.assertEqual(__a			,  [{'generated_text': ANY(__a   )}]   )
     self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'   )   )

     __UpperCamelCase       						= text_generator('This is a test'			,  return_full_text=__a   )
     self.assertEqual(__a			,  [{'generated_text': ANY(__a   )}]   )
     self.assertNotIn('This is a test'			,  outputs[0]['generated_text']   )

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

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

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

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

     with self.assertRaises(__a   ):
      __UpperCamelCase       						= text_generator('test'			,  return_full_text=__a			,  return_text=__a   )
     with self.assertRaises(__a   ):
      __UpperCamelCase       						= text_generator('test'			,  return_full_text=__a			,  return_tensors=__a   )
     with self.assertRaises(__a   ):
      __UpperCamelCase       						= text_generator('test'			,  return_text=__a			,  return_tensors=__a   )

     # 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__
     ):
      __UpperCamelCase       						= text_generator(''   )
      self.assertEqual(__a			,  [{'generated_text': ANY(__a   )}]   )
     else:
      with self.assertRaises((ValueError, AssertionError)   ):
       __UpperCamelCase       						= 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.
     __UpperCamelCase       						= ['RwkvForCausalLM', 'XGLMForCausalLM', 'GPTNeoXForCausalLM']
     if (
         tokenizer.model_max_length < 1_0000
         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   )

      __UpperCamelCase       						= text_generator('This is a test' * 500			,  handle_long_generation='hole'			,  max_new_tokens=20   )
      # Hole strategy cannot work
      with self.assertRaises(__a   ):
       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   ):

     '''simple docstring'''




     import torch

     # Classic `model_kwargs`
     __UpperCamelCase       						= 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   )
     __UpperCamelCase       						= pipe('This is a test'   )
     self.assertEqual(
         __a			,  [
             {
                 '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.)
     __UpperCamelCase       						= 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   )
     __UpperCamelCase       						= pipe('This is a test'   )
     self.assertEqual(
         __a			,  [
             {
                 '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
     __UpperCamelCase       						= 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   )
     __UpperCamelCase       						= pipe('This is a test'   )
     self.assertEqual(
         __a			,  [
             {
                 '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   ):

     '''simple docstring'''




     import torch

     __UpperCamelCase       						= 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   ):

     '''simple docstring'''




     import torch

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







    def 		UpperCAmelCase      (			self   ):

     '''simple docstring'''




     __UpperCamelCase       						= 'Hello world'
     __UpperCamelCase       						= pipeline('text-generation'			,  model='hf-internal-testing/tiny-random-gpt2'   )
     if text_generator.model.framework == "tf":
      __UpperCamelCase       						= logging.get_logger('transformers.generation.tf_utils'   )
     else:
      __UpperCamelCase       						= logging.get_logger('transformers.generation.utils'   )
     __UpperCamelCase       						= '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(__a   ) as cl:
      __UpperCamelCase       						= text_generator(__a			,  max_length=10			,  max_new_tokens=1   )
     self.assertIn(__a			,  cl.out   )

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

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





"""simple docstring"""

import argparse
from collections import defaultdict






def 							A		(       snake_case   :Tuple       ,      snake_case   :List[Any]       ,      snake_case   :List[Any]       ,      snake_case   :Union[str, Any]       ,      snake_case   :Optional[Any]	)   ->			Union[str, Any]:
 __UpperCamelCase       						= f'{file}_{class_name}_{test_name}'
 done_test[_id] += 1

 with open(snake_case       ,      'r'	) as f:
  __UpperCamelCase       						= f.readlines()

 __UpperCamelCase       						= f'class {class_name}('
 __UpperCamelCase       						= f'{4 * " "}def {test_name}('
 __UpperCamelCase       						= f'{8 * " "}{correct_line.split()[0]}'
 __UpperCamelCase       						= f'{1_6 * " "}{correct_line.split()[0]}'
 __UpperCamelCase       						= False
 __UpperCamelCase       						= False
 __UpperCamelCase       						= False
 __UpperCamelCase       						= False
 __UpperCamelCase       						= 0
 __UpperCamelCase       						= 0

 __UpperCamelCase       						= []
 for line in lines:
  if line.startswith(snake_case	):
   __UpperCamelCase       						= True
  elif in_class and line.startswith(snake_case	):
   __UpperCamelCase       						= True
  elif in_class and in_func and (line.startswith(snake_case	) or line.startswith(snake_case	)):
   __UpperCamelCase       						= len(line.split(correct_line.split()[0]	)[0]	)
   count += 1

   if count == done_test[_id]:
    __UpperCamelCase       						= True

  if in_class and in_func and in_line:
   if ")" not in line:
    continue
   else:
    __UpperCamelCase       						= True

  if in_class and in_func and in_line and insert_line:
   new_lines.append(f'{spaces * " "}{correct_line}'	)
   __UpperCamelCase       						= __UpperCamelCase       						= __UpperCamelCase       						= __UpperCamelCase       						= False
  else:
   new_lines.append(snake_case	)

 with open(snake_case       ,      'w'	) as f:
  for line in new_lines:
   f.write(snake_case	)






def 							A		(       snake_case   :List[str]       ,      snake_case   :str=None	)   ->			Tuple:
 if fail is not None:
  with open(snake_case       ,      'r'	) as f:
   __UpperCamelCase       						= {l.strip() for l in f.readlines()}
 else:
  __UpperCamelCase       						= None

 with open(snake_case       ,      'r'	) as f:
  __UpperCamelCase       						= f.readlines()

 __UpperCamelCase       						= defaultdict(snake_case	)
 for line in correct_lines:
  __UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase						,			__UpperCamelCase       						= line.split(';'	)
  if test_failures is None or "::".join([file, class_name, test_name]	) in test_failures:
   overwrite_file(snake_case       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case	)


if __name__ == "__main__":
      UpperCamelCase			:       str									=				argparse.ArgumentParser()
      parser.add_argument("--correct_filename", help="filename of tests with expected result")
      parser.add_argument("--fail_filename", help="filename of test failures", type=str, default=None)
      UpperCamelCase			:       List[Any]									=				parser.parse_args()

      main(args.correct_filename, args.fail_filename)



