




'''simple docstring'''



from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel





@require_tf
class     lowerCAmelCase   :
       __lowercase  :      Optional[Any]        =			MarianConfig
       __lowercase  :      Any        =			{}
       __lowercase  :      Tuple        =			'''gelu'''


       def __init__(   self    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE=13    ,  __SCREAMING_SNAKE_CASE=7    ,  __SCREAMING_SNAKE_CASE=True    ,  __SCREAMING_SNAKE_CASE=False    ,  __SCREAMING_SNAKE_CASE=99    ,  __SCREAMING_SNAKE_CASE=32    ,  __SCREAMING_SNAKE_CASE=2    ,  __SCREAMING_SNAKE_CASE=4    ,  __SCREAMING_SNAKE_CASE=37    ,  __SCREAMING_SNAKE_CASE=0.1    ,  __SCREAMING_SNAKE_CASE=0.1    ,  __SCREAMING_SNAKE_CASE=20    ,  __SCREAMING_SNAKE_CASE=2    ,  __SCREAMING_SNAKE_CASE=1    ,  __SCREAMING_SNAKE_CASE=0    ,  )      ->  List[Any]:



            '''simple docstring'''




            __snake_case    					=  parent
            __snake_case    					=  batch_size
            __snake_case    					=  seq_length
            __snake_case    					=  is_training
            __snake_case    					=  use_labels
            __snake_case    					=  vocab_size
            __snake_case    					=  hidden_size
            __snake_case    					=  num_hidden_layers
            __snake_case    					=  num_attention_heads
            __snake_case    					=  intermediate_size

            __snake_case    					=  hidden_dropout_prob
            __snake_case    					=  attention_probs_dropout_prob
            __snake_case    					=  max_position_embeddings
            __snake_case    					=  eos_token_id
            __snake_case    					=  pad_token_id
            __snake_case    					=  bos_token_id


       def       lowerCAmelCase      (   self  )      ->  List[Any]:



            '''simple docstring'''




            __snake_case    					=  ids_tensor([self.batch_size, self.seq_length - 1]    ,  self.vocab_size  )
            __snake_case    					=  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  )    ,  1  )
            __snake_case    					=  tf.concat([input_ids, eos_tensor]    ,  axis=1  )

            __snake_case    					=  ids_tensor([self.batch_size, self.seq_length]    ,  self.vocab_size  )

            __snake_case    					=  self.config_cls(
                vocab_size=self.vocab_size    ,  d_model=self.hidden_size    ,  encoder_layers=self.num_hidden_layers    ,  decoder_layers=self.num_hidden_layers    ,  encoder_attention_heads=self.num_attention_heads    ,  decoder_attention_heads=self.num_attention_heads    ,  encoder_ffn_dim=self.intermediate_size    ,  decoder_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_ids=[2]    ,  bos_token_id=self.bos_token_id    ,  pad_token_id=self.pad_token_id    ,  decoder_start_token_id=self.pad_token_id    ,  **self.config_updates    ,  )
            __snake_case    					=  prepare_marian_inputs_dict(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
            return config, inputs_dict




       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[int]:



            '''simple docstring'''




            __snake_case    					=  TFMarianModel(config=__SCREAMING_SNAKE_CASE  ).get_decoder()
            __snake_case    					=  inputs_dict['''input_ids''']

            __snake_case    					=  input_ids[:1, :]
            __snake_case    					=  inputs_dict['''attention_mask'''][:1, :]
            __snake_case    					=  inputs_dict['''head_mask''']
            __snake_case    					=  1

            # first forward pass
            __snake_case    					=  model(__SCREAMING_SNAKE_CASE    ,  attention_mask=__SCREAMING_SNAKE_CASE    ,  head_mask=__SCREAMING_SNAKE_CASE    ,  use_cache=__SCREAMING_SNAKE_CASE  )

            __snake_case					,    __snake_case    					=  outputs.to_tuple()

            # create hypothetical next token and extent to next_input_ids
            __snake_case    					=  ids_tensor((self.batch_size, 3)    ,  config.vocab_size  )
            __snake_case    					=  tf.cast(ids_tensor((self.batch_size, 3)    ,  2  )    ,  tf.inta  )

            # append to next input_ids and
            __snake_case    					=  tf.concat([input_ids, next_tokens]    ,  axis=-1  )
            __snake_case    					=  tf.concat([attention_mask, next_attn_mask]    ,  axis=-1  )

            __snake_case    					=  model(__SCREAMING_SNAKE_CASE    ,  attention_mask=__SCREAMING_SNAKE_CASE  )[0]
            __snake_case    					=  model(__SCREAMING_SNAKE_CASE    ,  attention_mask=__SCREAMING_SNAKE_CASE    ,  past_key_values=__SCREAMING_SNAKE_CASE  )[0]

            self.parent.assertEqual(next_tokens.shape[1]    ,  output_from_past.shape[1]  )

            # select random slice
            __snake_case    					=  int(ids_tensor((1,)    ,  output_from_past.shape[-1]  )  )
            __snake_case    					=  output_from_no_past[:, -3:, random_slice_idx]
            __snake_case    					=  output_from_past[:, :, random_slice_idx]

            # test that outputs are equal for slice
            tf.debugging.assert_near(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  rtol=1E-3  )







def 			_UpperCamelCase (_lowerCamelCase							:     Any   ,	_lowerCamelCase							:     Tuple   ,	_lowerCamelCase							:     int   ,	_lowerCamelCase							:     int=None   ,	_lowerCamelCase							:     Dict=None   ,	_lowerCamelCase							:     List[str]=None   ,	_lowerCamelCase							:     Tuple=None   ,	_lowerCamelCase							:     Dict=None   ,	)->   str:




     '''simple docstring'''





     if attention_mask is None:
          __snake_case    					=  tf.cast(tf.math.not_equal(_lowerCamelCase   ,	config.pad_token_id	)   ,	tf.inta	)
     if decoder_attention_mask is None:
          __snake_case    					=  tf.concat(
              [
                  tf.ones(decoder_input_ids[:, :1].shape   ,	dtype=tf.inta	),
                  tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]   ,	config.pad_token_id	)   ,	tf.inta	),
              ]   ,	axis=-1   ,	)
     if head_mask is None:
          __snake_case    					=  tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
     if decoder_head_mask is None:
          __snake_case    					=  tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
     if cross_attn_head_mask is None:
          __snake_case    					=  tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
     return {
         "input_ids": input_ids,
         "decoder_input_ids": decoder_input_ids,
         "attention_mask": attention_mask,
         "decoder_attention_mask": decoder_attention_mask,
         "head_mask": head_mask,
         "decoder_head_mask": decoder_head_mask,
         "cross_attn_head_mask": cross_attn_head_mask,
     }





@require_tf
class     lowerCAmelCase   ( __lowerCAmelCase				,   __lowerCAmelCase				,   unittest.TestCase):
       __lowercase  :      Dict        =			(TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
       __lowercase  :      Dict        =			(TFMarianMTModel,) if is_tf_available() else ()
       __lowercase  :      Optional[int]        =			(
           {
               '''conversational''': TFMarianMTModel,
               '''feature-extraction''': TFMarianModel,
               '''summarization''': TFMarianMTModel,
               '''text2text-generation''': TFMarianMTModel,
               '''translation''': TFMarianMTModel,
           }
           if is_tf_available()
           else {}
       )
       __lowercase  :      Tuple        =			True
       __lowercase  :      Dict        =			False
       __lowercase  :      str        =			False


       def       lowerCAmelCase      (   self  )      ->  Union[str, Any]:



            '''simple docstring'''




            __snake_case    					=  TFMarianModelTester(self  )
            __snake_case    					=  ConfigTester(self    ,  config_class=__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self  )      ->  Optional[Any]:



            '''simple docstring'''




            self.config_tester.run_common_tests()




       def       lowerCAmelCase      (   self  )      ->  Optional[int]:



            '''simple docstring'''




            __snake_case    					=  self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_decoder_model_past_large_inputs(*__SCREAMING_SNAKE_CASE  )





@require_tf
class     lowerCAmelCase   ( unittest.TestCase):
       __lowercase  :      str        =			1000  # show more chars for failing integration tests


       @classmethod
       def       lowerCAmelCase      (   cls  )      ->  None:



            '''simple docstring'''




            __snake_case    					=  F'''Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}'''
            return cls


       @cached_property
       def       lowerCAmelCase      (   self  )      ->  MarianTokenizer:



            '''simple docstring'''




            return AutoTokenizer.from_pretrained(self.model_name  )


       @property
       def       lowerCAmelCase      (   self  )      ->  int:



            '''simple docstring'''




            return self.tokenizer.eos_token_id


       @cached_property
       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            warnings.simplefilter('''error'''  )
            __snake_case    					=  TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name  )
            assert isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
            __snake_case    					=  model.config
            self.assertListEqual(c.bad_words_ids    ,  [[c.pad_token_id]]  )
            self.assertEqual(c.max_length    ,  512  )
            self.assertEqual(c.decoder_start_token_id    ,  c.pad_token_id  )
            return model


       def       lowerCAmelCase      (   self    ,  **__SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            __snake_case    					=  self.translate_src_text(**__SCREAMING_SNAKE_CASE  )
            self.assertListEqual(self.expected_text    ,  __SCREAMING_SNAKE_CASE  )




       def       lowerCAmelCase      (   self    ,  **__SCREAMING_SNAKE_CASE  )      ->  Optional[int]:



            '''simple docstring'''




            __snake_case    					=  self.tokenizer(self.src_text    ,  **__SCREAMING_SNAKE_CASE    ,  padding=__SCREAMING_SNAKE_CASE    ,  return_tensors='''tf'''  )
            __snake_case    					=  self.model.generate(
                model_inputs.input_ids    ,  attention_mask=model_inputs.attention_mask    ,  num_beams=2    ,  max_length=128  )
            __snake_case    					=  self.tokenizer.batch_decode(generated_ids.numpy()    ,  skip_special_tokens=__SCREAMING_SNAKE_CASE  )
            return generated_words





@require_sentencepiece
@require_tokenizers
@require_tf
class     lowerCAmelCase   ( __lowerCAmelCase):
       __lowercase  :      Dict        =			'''mt'''
       __lowercase  :      int        =			'''en'''
       __lowercase  :      Optional[Any]        =			['''Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.''']
       __lowercase  :      Any        =			['''Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.''']


       @unittest.skip('''Skipping until #12647 is resolved.'''  )
       @slow
       def       lowerCAmelCase      (   self  )      ->  int:



            '''simple docstring'''




            self._assert_generated_batch_equal_expected()





@require_sentencepiece
@require_tokenizers
@require_tf
class     lowerCAmelCase   ( __lowerCAmelCase):
       __lowercase  :      Union[str, Any]        =			'''en'''
       __lowercase  :      List[Any]        =			'''zh'''
       __lowercase  :      Dict        =			['''My name is Wolfgang and I live in Berlin''']
       __lowercase  :      Any        =			['''我叫沃尔夫冈 我住在柏林''']


       @unittest.skip('''Skipping until #12647 is resolved.'''  )
       @slow
       def       lowerCAmelCase      (   self  )      ->  Optional[int]:



            '''simple docstring'''




            self._assert_generated_batch_equal_expected()





@require_sentencepiece
@require_tokenizers
@require_tf
class     lowerCAmelCase   ( __lowerCAmelCase):
       __lowercase  :      int        =			'''en'''
       __lowercase  :      List[str]        =			'''ROMANCE'''
       __lowercase  :      Union[str, Any]        =			[
           '''>>fr<< Don\'t spend so much time watching TV.''',
           '''>>pt<< Your message has been sent.''',
           '''>>es<< He\'s two years older than me.''',
       ]
       __lowercase  :      Union[str, Any]        =			[
           '''Ne passez pas autant de temps à regarder la télé.''',
           '''A sua mensagem foi enviada.''',
           '''Es dos años más viejo que yo.''',
       ]


       @unittest.skip('''Skipping until #12647 is resolved.'''  )
       @slow
       def       lowerCAmelCase      (   self  )      ->  Tuple:



            '''simple docstring'''




            self._assert_generated_batch_equal_expected()




       @unittest.skip('''Skipping until #12647 is resolved.'''  )
       @slow
       def       lowerCAmelCase      (   self  )      ->  str:



            '''simple docstring'''




            __snake_case    					=  TranslationPipeline(self.model    ,  self.tokenizer    ,  framework='''tf'''  )
            __snake_case    					=  pipeline(self.src_text  )
            self.assertEqual(self.expected_text    ,  [x['''translation_text'''] for x in output]  )












'''simple docstring'''



import itertools
from dataclasses import dataclass
from io import StringIO
from typing import Optional

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


UpperCAmelCase_       :  str      =		datasets.utils.logging.get_logger(__name__)





@dataclass
class     lowerCAmelCase   ( datasets.BuilderConfig):
       __lowercase  :      Optional[datasets.Features]        =			None
       __lowercase  :      str        =			"utf-8"
       __lowercase  :      Optional[str]        =			None
       __lowercase  :      int        =			10 << 20  # 10MB
       __lowercase  :      bool        =			False
       __lowercase  :      str        =			"line"





class     lowerCAmelCase   ( datasets.ArrowBasedBuilder):
       __lowercase  :      Dict        =			TextConfig


       def       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




            return datasets.DatasetInfo(features=self.config.features  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            if not self.config.data_files:
                 raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''  )
            __snake_case    					=  dl_manager.download_and_extract(self.config.data_files  )
            if isinstance(__SCREAMING_SNAKE_CASE    ,  (str, list, tuple)  ):
                 __snake_case    					=  data_files
                 if isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ):
                      __snake_case    					=  [files]
                 __snake_case    					=  [dl_manager.iter_files(__SCREAMING_SNAKE_CASE  ) for file in files]
                 return [datasets.SplitGenerator(name=datasets.Split.TRAIN    ,  gen_kwargs={'''files''': files}  )]
            __snake_case    					=  []
            for split_name, files in data_files.items():
                 if isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ):
                      __snake_case    					=  [files]
                 __snake_case    					=  [dl_manager.iter_files(__SCREAMING_SNAKE_CASE  ) for file in files]
                 splits.append(datasets.SplitGenerator(name=__SCREAMING_SNAKE_CASE    ,  gen_kwargs={'''files''': files}  )  )
            return splits


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  pa.Table:



            '''simple docstring'''




            if self.config.features is not None:
                 __snake_case    					=  self.config.features.arrow_schema
                 if all(not require_storage_cast(__SCREAMING_SNAKE_CASE  ) for feature in self.config.features.values()  ):
                      # cheaper cast
                      __snake_case    					=  pa_table.cast(__SCREAMING_SNAKE_CASE  )
                 else:
                      # more expensive cast; allows str <-> int/float or str to Audio for example
                      __snake_case    					=  table_cast(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
                 return pa_table
            else:
                 return pa_table.cast(pa.schema({'''text''': pa.string()}  )  )




       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            __snake_case    					=  list(self.config.features  ) if self.config.features is not None else ['''text''']
            for file_idx, file in enumerate(itertools.chain.from_iterable(__SCREAMING_SNAKE_CASE  )  ):
                 # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
                 with open(__SCREAMING_SNAKE_CASE    ,  encoding=self.config.encoding    ,  errors=self.config.errors  ) as f:
                      if self.config.sample_by == "line":
                           __snake_case    					=  0
                           while True:
                                __snake_case    					=  f.read(self.config.chunksize  )
                                if not batch:
                                     break
                                batch += f.readline()  # finish current line
                                # StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
                                __snake_case    					=  StringIO(__SCREAMING_SNAKE_CASE  ).readlines()
                                if not self.config.keep_linebreaks:
                                     __snake_case    					=  [line.rstrip('''\n'''  ) for line in batch]
                                __snake_case    					=  pa.Table.from_arrays([pa.array(__SCREAMING_SNAKE_CASE  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                # Uncomment for debugging (will print the Arrow table size and elements)
                                # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                                batch_idx += 1
                      elif self.config.sample_by == "paragraph":
                           __snake_case    					=  0
                           __snake_case    					=  ''''''
                           while True:
                                __snake_case    					=  f.read(self.config.chunksize  )
                                if not new_batch:
                                     break
                                batch += new_batch
                                batch += f.readline()  # finish current line
                                __snake_case    					=  batch.split('''\n\n'''  )
                                __snake_case    					=  pa.Table.from_arrays(
                                    [pa.array([example for example in batch[:-1] if example]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                # Uncomment for debugging (will print the Arrow table size and elements)
                                # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                                batch_idx += 1
                                __snake_case    					=  batch[-1]
                           if batch:
                                __snake_case    					=  pa.Table.from_arrays([pa.array([batch]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                      elif self.config.sample_by == "document":
                           __snake_case    					=  f.read()
                           __snake_case    					=  pa.Table.from_arrays([pa.array([text]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                           yield file_idx, self._cast_table(__SCREAMING_SNAKE_CASE  )






