





'''simple docstring'''







import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


A						    =							'''\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
'''

A						    =							'''
This metric wrap the official scoring script for version 2 of the Stanford Question
Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.

SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions
written adversarially by crowdworkers to look similar to answerable ones.
To do well on SQuAD2.0, systems must not only answer questions when possible, but also
determine when no answer is supported by the paragraph and abstain from answering.
'''

A						    =							'''
Computes SQuAD v2 scores (F1 and EM).
Args:
    predictions: List of triple for question-answers to score with the following elements:
        - the question-answer \'id\' field as given in the references (see below)
        - the text of the answer
        - the probability that the question has no answer
    references: List of question-answers dictionaries with the following key-values:
            - \'id\': id of the question-answer pair (see above),
            - \'answers\': a list of Dict {\'text\': text of the answer as a string}
    no_answer_threshold: float
        Probability threshold to decide that a question has no answer.
Returns:
    \'exact\': Exact match (the normalized answer exactly match the gold answer)
    \'f1\': The F-score of predicted tokens versus the gold answer
    \'total\': Number of score considered
    \'HasAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'HasAns_f1\': The F-score of predicted tokens versus the gold answer
    \'HasAns_total\': Number of score considered
    \'NoAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'NoAns_f1\': The F-score of predicted tokens versus the gold answer
    \'NoAns_total\': Number of score considered
    \'best_exact\': Best exact match (with varying threshold)
    \'best_exact_thresh\': No-answer probability threshold associated to the best exact match
    \'best_f1\': Best F1 (with varying threshold)
    \'best_f1_thresh\': No-answer probability threshold associated to the best F1
Examples:

    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\', \'no_answer_probability\': 0.}]
    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> squad_v2_metric = datasets.load_metric("squad_v2")
    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact\': 100.0, \'f1\': 100.0, \'total\': 1, \'HasAns_exact\': 100.0, \'HasAns_f1\': 100.0, \'HasAns_total\': 1, \'best_exact\': 100.0, \'best_exact_thresh\': 0.0, \'best_f1\': 100.0, \'best_f1_thresh\': 0.0}
'''



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,   _KWARGS_DESCRIPTION      )
class 							__SCREAMING_SNAKE_CASE			(       datasets.Metric      ):
      '''simple docstring'''



      def   _lowerCamelCase		(							self    : Tuple  )		->  Dict:
           return datasets.MetricInfo(
               description=_DESCRIPTION   ,citation=_CITATION   ,inputs_description=_KWARGS_DESCRIPTION   ,features=datasets.Features(
                   {
                       'predictions': {
                           'id': datasets.Value('string'  ),
                           'prediction_text': datasets.Value('string'  ),
                           'no_answer_probability': datasets.Value('float32'  ),
                       },
                       'references': {
                           'id': datasets.Value('string'  ),
                           'answers': datasets.features.Sequence(
                               {'text': datasets.Value('string'  ), 'answer_start': datasets.Value('int32'  )}  ),
                       },
                   }  )   ,codebase_urls=['https://rajpurkar.github.io/SQuAD-explorer/']   ,reference_urls=['https://rajpurkar.github.io/SQuAD-explorer/']   ,)



      def   _lowerCamelCase		(							self    : Optional[int]   ,UpperCamelCase    : Optional[Any]   ,UpperCamelCase    : Any   ,UpperCamelCase    : int=1.0  )		->  Any:
           _lowercase							:      Optional[int]  						=							{p['id']: p['no_answer_probability'] for p in predictions}
           _lowercase							:      Union[str, Any]  						=							[{'paragraphs': [{'qas': references}]}]
           _lowercase							:      str  						=							{p['id']: p['prediction_text'] for p in predictions}

           _lowercase							:      Tuple  						=							make_qid_to_has_ans(UpperCamelCase  )  # maps qid to True/False
           _lowercase							:      Dict  						=							[k for k, v in qid_to_has_ans.items() if v]
           _lowercase							:      Tuple  						=							[k for k, v in qid_to_has_ans.items() if not v]

           _lowercase						, _lowercase							:      Tuple  						=							get_raw_scores(UpperCamelCase   ,UpperCamelCase  )
           _lowercase							:      Tuple  						=							apply_no_ans_threshold(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							apply_no_ans_threshold(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )
           _lowercase							:      Optional[int]  						=							make_eval_dict(UpperCamelCase   ,UpperCamelCase  )

           if has_ans_qids:
                _lowercase							:      Dict  						=							make_eval_dict(UpperCamelCase   ,UpperCamelCase   ,qid_list=UpperCamelCase  )
                merge_eval(UpperCamelCase   ,UpperCamelCase   ,'HasAns'  )
           if no_ans_qids:
                _lowercase							:      Tuple  						=							make_eval_dict(UpperCamelCase   ,UpperCamelCase   ,qid_list=UpperCamelCase  )
                merge_eval(UpperCamelCase   ,UpperCamelCase   ,'NoAns'  )
           find_all_best_thresh(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )
           return dict(UpperCamelCase  )






'''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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[int]       ,  lowerCAmelCase__     :	Optional[int]       ,  lowerCAmelCase__     :	int=None       ,  lowerCAmelCase__     :	List[Any]=None)			->						Union[str, Any]:

     '''simple docstring'''




     if attention_mask is None:
          _lowercase							:      str  						=							np.where(input_ids != config.pad_token_id       ,  1       ,  0)
     return {
         "input_ids": input_ids,
         "attention_mask": attention_mask,
     }



@require_flax
class 							__SCREAMING_SNAKE_CASE			:
      '''simple docstring'''



      def __init__(							self    : List[str]   ,UpperCamelCase    : Dict   ,UpperCamelCase    : Union[str, Any]=13   ,UpperCamelCase    : Optional[Any]=7   ,UpperCamelCase    : Optional[int]=True   ,UpperCamelCase    : Any=False   ,UpperCamelCase    : int=99   ,UpperCamelCase    : Optional[int]=16   ,UpperCamelCase    : Optional[int]=2   ,UpperCamelCase    : int=4   ,UpperCamelCase    : Optional[int]=4   ,UpperCamelCase    : Dict="gelu"   ,UpperCamelCase    : Optional[Any]=0.1   ,UpperCamelCase    : Optional[int]=0.1   ,UpperCamelCase    : str=20   ,UpperCamelCase    : List[Any]=2   ,UpperCamelCase    : Union[str, Any]=1   ,UpperCamelCase    : List[Any]=0   ,UpperCamelCase    : Optional[Any]=16   ,UpperCamelCase    : List[Any]=16   ,UpperCamelCase    : Optional[Any]=0.0_2   ,)		->  Dict:
           _lowercase							:      List[Any]  						=							parent
           _lowercase							:      Union[str, Any]  						=							batch_size
           _lowercase							:      int  						=							seq_length
           _lowercase							:      Dict  						=							is_training
           _lowercase							:      Tuple  						=							use_labels
           _lowercase							:      List[str]  						=							vocab_size
           _lowercase							:      str  						=							hidden_size
           _lowercase							:      List[str]  						=							num_hidden_layers
           _lowercase							:      List[str]  						=							num_attention_heads
           _lowercase							:      Dict  						=							intermediate_size
           _lowercase							:      Tuple  						=							hidden_act
           _lowercase							:      List[str]  						=							hidden_dropout_prob
           _lowercase							:      Optional[Any]  						=							attention_probs_dropout_prob
           _lowercase							:      List[str]  						=							max_position_embeddings
           _lowercase							:      Union[str, Any]  						=							eos_token_id
           _lowercase							:      Tuple  						=							pad_token_id
           _lowercase							:      str  						=							bos_token_id
           _lowercase							:      List[str]  						=							embed_dim
           _lowercase							:      Tuple  						=							word_embed_proj_dim
           _lowercase							:      Tuple  						=							initializer_range
           _lowercase							:      Any  						=							False

      def   _lowerCamelCase		(							self    : str  )		->  Union[str, Any]:
           _lowercase							:      str  						=							np.clip(ids_tensor([self.batch_size, self.seq_length - 1]   ,self.vocab_size  )   ,3   ,self.vocab_size  )
           _lowercase							:      int  						=							np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)   ,dtype=np.intaa  ))   ,-1  )

           _lowercase							:      Optional[int]  						=							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=UpperCamelCase   ,word_embed_proj_dim=self.word_embed_proj_dim   ,initializer_range=self.initializer_range   ,use_cache=UpperCamelCase   ,)
           _lowercase							:      int  						=							prepare_opt_inputs_dict(UpperCamelCase   ,UpperCamelCase  )
           return config, inputs_dict

      def   _lowerCamelCase		(							self    : Tuple  )		->  List[str]:
           _lowercase						, _lowercase							:      Tuple  						=							self.prepare_config_and_inputs()
           return config, inputs_dict

      def   _lowerCamelCase		(							self    : Union[str, Any]   ,UpperCamelCase    : str   ,UpperCamelCase    : List[Any]   ,UpperCamelCase    : Optional[Any]  )		->  Optional[Any]:
           _lowercase							:      Dict  						=							20
           _lowercase							:      Any  						=							model_class_name(UpperCamelCase  )

           _lowercase							:      Dict  						=							inputs_dict['input_ids']
           _lowercase							:      Dict  						=							inputs_dict['attention_mask']

           _lowercase							:      Dict  						=							model.init_cache(input_ids.shape[0]   ,UpperCamelCase  )
           _lowercase							:      List[str]  						=							jnp.ones((input_ids.shape[0], max_length)   ,dtype='i4'  )

           _lowercase							:      Dict  						=							jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1  )[None, :]   ,(input_ids.shape[0], input_ids.shape[-1] - 1)   ,)
           _lowercase							:      List[Any]  						=							model(
               input_ids[:, :-1]   ,attention_mask=UpperCamelCase   ,past_key_values=UpperCamelCase   ,position_ids=UpperCamelCase   ,)

           _lowercase							:      str  						=							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,dtype='i4'  )
           _lowercase							:      Optional[int]  						=							model(
               input_ids[:, -1:]   ,attention_mask=UpperCamelCase   ,past_key_values=outputs_cache.past_key_values   ,position_ids=UpperCamelCase   ,)

           _lowercase							:      Dict  						=							model(UpperCamelCase  )

           _lowercase							:      Dict  						=							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   _lowerCamelCase		(							self    : List[Any]   ,UpperCamelCase    : List[Any]   ,UpperCamelCase    : List[Any]   ,UpperCamelCase    : str  )		->  Optional[Any]:
           _lowercase							:      Optional[int]  						=							20
           _lowercase							:      Dict  						=							model_class_name(UpperCamelCase  )

           _lowercase						, _lowercase							:      Tuple  						=							(
               inputs_dict['input_ids'],
               inputs_dict['attention_mask'],
           )

           _lowercase							:      int  						=							jnp.concatenate(
               [
                   attention_mask,
                   jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])  ),
               ]   ,axis=-1   ,)

           _lowercase							:      Dict  						=							model.init_cache(input_ids.shape[0]   ,UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1  )[None, :]   ,(input_ids.shape[0], input_ids.shape[-1] - 1)   ,)

           _lowercase							:      List[Any]  						=							model(
               input_ids[:, :-1]   ,attention_mask=UpperCamelCase   ,past_key_values=UpperCamelCase   ,position_ids=UpperCamelCase   ,)
           _lowercase							:      Any  						=							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,dtype='i4'  )
           _lowercase							:      List[Any]  						=							model(
               input_ids[:, -1:]   ,past_key_values=outputs_cache.past_key_values   ,attention_mask=UpperCamelCase   ,position_ids=UpperCamelCase   ,)

           _lowercase							:      Tuple  						=							model(UpperCamelCase   ,attention_mask=UpperCamelCase  )

           _lowercase							:      List[Any]  						=							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 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__							,   unittest.TestCase							,   lowerCAmelCase__      ):
      '''simple docstring'''



      lowerCAmelCase__       :      List[str] 		=		(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
      lowerCAmelCase__       :      List[Any] 		=		() if is_flax_available() else ()

      def   _lowerCamelCase		(							self    : Optional[int]  )		->  str:
           _lowercase							:      Dict  						=							FlaxOPTModelTester(self  )

      def   _lowerCamelCase		(							self    : int  )		->  List[str]:
           _lowercase						, _lowercase							:      List[Any]  						=							self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                self.model_tester.check_use_cache_forward(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )

      def   _lowerCamelCase		(							self    : Any  )		->  Dict:
           _lowercase						, _lowercase							:      Union[str, 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(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )



      @slow
      def   _lowerCamelCase		(							self    : Optional[int]  )		->  int:
           for model_class_name in self.all_model_classes:
                _lowercase							:      Optional[int]  						=							model_class_name.from_pretrained('facebook/opt-125m'  )
                _lowercase							:      Any  						=							np.ones((1, 1)  ) * model.config.eos_token_id
                _lowercase							:      str  						=							model(UpperCamelCase  )
                self.assertIsNotNone(UpperCamelCase  )



@require_sentencepiece
@require_flax
class 							__SCREAMING_SNAKE_CASE			(       unittest.TestCase      ):
      '''simple docstring'''



      @slow
      def   _lowerCamelCase		(							self    : int  )		->  Optional[int]:
           _lowercase							:      Optional[int]  						=							FlaxOPTModel.from_pretrained('facebook/opt-350m'  )
           _lowercase							:      int  						=							jnp.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]  )
           _lowercase							:      Tuple  						=							model(input_ids=UpperCamelCase  ).last_hidden_state
           _lowercase							:      Union[str, Any]  						=							(1, 11, 512)
           self.assertEqual(output.shape   ,UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							jnp.array(
               [[-0.2_8_6_7, -1.9_2_5_6, -0.3_0_6_2], [-1.2_7_1_1, -0.1_3_3_7, -0.1_8_9_7], [0.4_1_0_9, 0.1_1_8_7, -1.3_1_4_2]]  )
           self.assertTrue(jnp.allclose(output[:, :3, :3]   ,UpperCamelCase   ,atol=4e-2  )  )



@require_flax
@slow
class 							__SCREAMING_SNAKE_CASE			(       unittest.TestCase      ):
      '''simple docstring'''



      def   _lowerCamelCase		(							self    : Optional[Any]  )		->  Tuple:
           super().setUp()
           _lowercase							:      Optional[int]  						=							'facebook/opt-350m'



      def   _lowerCamelCase		(							self    : Optional[int]  )		->  Union[str, Any]:
           _lowercase							:      Tuple  						=							FlaxOPTForCausalLM.from_pretrained(self.path_model  )
           _lowercase							:      str  						=							GPTaTokenizer.from_pretrained(self.path_model  )

           _lowercase							:      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
           _lowercase							:      str  						=							tokenizer(UpperCamelCase   ,return_tensors='jax'   ,padding=UpperCamelCase   ,add_special_tokens=UpperCamelCase  )
           _lowercase							:      Optional[int]  						=							model(inputs.input_ids   ,attention_mask=inputs.attention_mask  )[0].mean(axis=-1  )
           _lowercase							:      int  						=							jnp.array(
               [
                   [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
                   [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
                   [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
                   [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
               ]  )
           self.assertTrue(jnp.allclose(UpperCamelCase   ,UpperCamelCase   ,atol=4e-2  )  )

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



@require_flax
@slow
class 							__SCREAMING_SNAKE_CASE			(       unittest.TestCase      ):
      '''simple docstring'''



      @property
      def   _lowerCamelCase		(							self    : Optional[int]  )		->  Optional[int]:
           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   _lowerCamelCase		(							self    : int  )		->  Optional[int]:
           _lowercase							:      Dict  						=							'facebook/opt-125m'

           _lowercase							:      Optional[int]  						=							[
               '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',
           ]

           _lowercase							:      int  						=							[]

           _lowercase							:      Tuple  						=							FlaxOPTForCausalLM.from_pretrained(UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							GPTaTokenizer.from_pretrained(UpperCamelCase  )

           for prompt in self.prompts:
                _lowercase							:      List[str]  						=							tokenizer(UpperCamelCase   ,return_tensors='jax'  ).input_ids

                _lowercase							:      Optional[Any]  						=							model.generate(UpperCamelCase   ,max_length=10  )
                _lowercase							:      Union[str, Any]  						=							generated_ids[0]

                _lowercase							:      Any  						=							tokenizer.batch_decode(UpperCamelCase   ,skip_special_tokens=UpperCamelCase  )
                predicted_outputs += generated_string

           self.assertListEqual(UpperCamelCase   ,UpperCamelCase  )

      def   _lowerCamelCase		(							self    : List[str]  )		->  Dict:
           _lowercase							:      int  						=							'facebook/opt-350m'

           _lowercase							:      Tuple  						=							[
               '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',
           ]

           _lowercase							:      Union[str, Any]  						=							[]
           _lowercase							:      List[Any]  						=							FlaxOPTForCausalLM.from_pretrained(UpperCamelCase  )
           _lowercase							:      Optional[int]  						=							GPTaTokenizer.from_pretrained(UpperCamelCase  )

           for prompt in self.prompts:
                _lowercase							:      Optional[int]  						=							tokenizer(UpperCamelCase   ,return_tensors='jax'  ).input_ids

                _lowercase							:      Tuple  						=							model.generate(UpperCamelCase   ,max_length=10  )
                _lowercase							:      Any  						=							generated_ids[0]

                _lowercase							:      Tuple  						=							tokenizer.batch_decode(UpperCamelCase   ,skip_special_tokens=UpperCamelCase  )
                predicted_outputs += generated_string

           self.assertListEqual(UpperCamelCase   ,UpperCamelCase  )

      def   _lowerCamelCase		(							self    : Tuple  )		->  Union[str, Any]:
           _lowercase							:      Union[str, Any]  						=							'facebook/opt-125m'
           _lowercase							:      List[Any]  						=							[
               'Today is a beautiful day and I want to thank',
               'In the city of Rome Canaver Canaver Canaver Canaver',
           ]
           _lowercase							:      List[str]  						=							FlaxOPTForCausalLM.from_pretrained(UpperCamelCase  )
           _lowercase							:      Optional[int]  						=							GPTaTokenizer.from_pretrained(UpperCamelCase  )
           _lowercase							:      Optional[Any]  						=							tokenizer(
               [
                   'Today is a beautiful day and I want to',
                   'In the city of',
               ]   ,return_tensors='jax'   ,padding=UpperCamelCase   ,)

           _lowercase							:      str  						=							jax.jit(model.generate  )

           _lowercase							:      List[str]  						=							jit_generate(inputs['input_ids']   ,attention_mask=inputs['attention_mask']  ).sequences

           _lowercase							:      Union[str, Any]  						=							tokenizer.batch_decode(UpperCamelCase   ,skip_special_tokens=UpperCamelCase  )

           self.assertIsNotNone(UpperCamelCase   ,UpperCamelCase  )



      def   _lowerCamelCase		(							self    : Dict  )		->  List[str]:
           _lowercase							:      str  						=							'facebook/opt-350m'

           _lowercase							:      List[str]  						=							GPTaTokenizer.from_pretrained(UpperCamelCase  )
           _lowercase							:      List[str]  						=							FlaxOPTForCausalLM.from_pretrained(UpperCamelCase  )

           _lowercase							:      int  						=							'left'

           # use different length sentences to test batching
           _lowercase							:      Tuple  						=							[
               'Hello, my dog is a little',
               'Today, I',
           ]

           _lowercase							:      int  						=							tokenizer(UpperCamelCase   ,return_tensors='jax'   ,padding=UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							inputs['input_ids']

           _lowercase							:      str  						=							model.generate(input_ids=UpperCamelCase   ,attention_mask=inputs['attention_mask']   ,trace=UpperCamelCase  )

           _lowercase							:      Tuple  						=							tokenizer(sentences[0]   ,return_tensors='jax'  ).input_ids
           _lowercase							:      Optional[Any]  						=							model.generate(input_ids=UpperCamelCase  )

           _lowercase							:      Optional[int]  						=							inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].sum()
           _lowercase							:      Optional[int]  						=							tokenizer(sentences[1]   ,return_tensors='jax'  ).input_ids
           _lowercase							:      List[str]  						=							model.generate(input_ids=UpperCamelCase   ,max_length=model.config.max_length - num_paddings  )

           _lowercase							:      Dict  						=							tokenizer.batch_decode(outputs[0]   ,skip_special_tokens=UpperCamelCase  )
           _lowercase							:      Optional[Any]  						=							tokenizer.decode(output_non_padded[0][0]   ,skip_special_tokens=UpperCamelCase  )
           _lowercase							:      Optional[Any]  						=							tokenizer.decode(output_padded[0][0]   ,skip_special_tokens=UpperCamelCase  )

           _lowercase							:      Union[str, Any]  						=							[
               '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(UpperCamelCase   ,UpperCamelCase  )
           self.assertListEqual(UpperCamelCase   ,[non_padded_sentence, padded_sentence]  )