

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   lowercase_   (  unittest.TestCase					):
      UpperCamelCase_      :		Tuple										=    MODEL_FOR_CAUSAL_LM_MAPPING
      UpperCamelCase_      :		str										=    TF_MODEL_FOR_CAUSAL_LM_MAPPING



      @require_torch
      def 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					Dict:
        _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=A__		)
        self.assertEqual(
            A__      , [
                {
                    '''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(
            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@@'''
                        )
                    }
                ],
            ]      , )

        _snake_case						=							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__		)},
            ]      , )
        _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=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				:    Union[str, Any]		)   ->					Tuple:
        _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=A__		)
        self.assertEqual(
            A__      , [
                {
                    '''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=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				:    Tuple      , A__				:    List[str]      , A__				:    Dict      , A__				:    Optional[int]		)   ->					Optional[int]:
        _snake_case						=							TextGenerationPipeline(model=A__      , tokenizer=A__		)
        return text_generator, ["This is a test", "Another test"]



      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					int:
        _snake_case						=							'''Hello I believe in'''
        _snake_case						=							pipeline('''text-generation'''      , model='''hf-internal-testing/tiny-random-gpt2'''		)
        _snake_case						=							text_generator(A__		)
        self.assertEqual(
            A__      , [{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]      , )

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



      def 						UpperCamelCase_		(      self				:    Tuple      , A__				:    Dict      , A__				:    Optional[int]		)   ->					Union[str, Any]:
        _snake_case						=							text_generator.model
        _snake_case						=							text_generator.tokenizer

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

        _snake_case						=							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''']		)

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

        _snake_case						=							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'''		)		)

        _snake_case						=							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:
          _snake_case						=							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__		):
          _snake_case						=							text_generator('''test'''      , return_full_text=A__      , return_text=A__		)
        with self.assertRaises(A__		):
          _snake_case						=							text_generator('''test'''      , return_full_text=A__      , return_tensors=A__		)
        with self.assertRaises(A__		):
          _snake_case						=							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__
        ):
          _snake_case						=							text_generator(''''''		)
          self.assertEqual(A__      , [{'''generated_text''': ANY(A__		)}]		)
        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 < 10000
            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(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				:    str		)   ->					Tuple:
        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(
            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.)
        _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(
            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
        _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(
            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				:    List[str]		)   ->					List[Any]:
        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				:    List[Any]		)   ->					Tuple:
        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=A__      , top_p=0.5		)







      def 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					Tuple:
        _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(A__		) as cl:
          _snake_case						=							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:
          _snake_case						=							text_generator(A__      , max_new_tokens=1		)
        self.assertNotIn(A__      , cl.out		)

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




# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__A      		=     {'''configuration_mra''': ['''MRA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MraConfig''']}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A      		=     [
     '''MRA_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''MraForMaskedLM''',
     '''MraForMultipleChoice''',
     '''MraForQuestionAnswering''',
     '''MraForSequenceClassification''',
     '''MraForTokenClassification''',
     '''MraLayer''',
     '''MraModel''',
     '''MraPreTrainedModel''',
 ]


if TYPE_CHECKING:
 from .configuration_mra import MRA_PRETRAINED_CONFIG_ARCHIVE_MAP, MraConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_mra import (
      MRA_PRETRAINED_MODEL_ARCHIVE_LIST,
      MraForMaskedLM,
      MraForMultipleChoice,
      MraForQuestionAnswering,
      MraForSequenceClassification,
      MraForTokenClassification,
      MraLayer,
      MraModel,
      MraPreTrainedModel,
  )
else:
 import sys

 __A      		=     _LazyModule(__name__, globals()['''__file__'''], _import_structure)

