


from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available
from transformers.testing_utils import 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, TFBlenderbotForConditionalGeneration, TFBlenderbotModel







@require_tf
class 	__lowercase    :





							"""simple docstring"""







							_UpperCAmelCase      =						BlenderbotConfig
							_UpperCAmelCase      =						{}
							_UpperCAmelCase      =						"""gelu"""






							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__=1_3			,						lowerCAmelCase__=7			,						lowerCAmelCase__=True			,						lowerCAmelCase__=False			,						lowerCAmelCase__=9_9			,						lowerCAmelCase__=3_2			,						lowerCAmelCase__=2			,						lowerCAmelCase__=4			,						lowerCAmelCase__=3_7			,						lowerCAmelCase__=0.1			,						lowerCAmelCase__=0.1			,						lowerCAmelCase__=2_0			,						lowerCAmelCase__=2			,						lowerCAmelCase__=1			,						lowerCAmelCase__=0			,						):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					parent
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					batch_size
										SCREAMING_SNAKE_CASE_		:     Dict					=					seq_length
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					is_training
										SCREAMING_SNAKE_CASE_		:     Dict					=					use_labels
										SCREAMING_SNAKE_CASE_		:     Any					=					vocab_size
										SCREAMING_SNAKE_CASE_		:     int					=					hidden_size
										SCREAMING_SNAKE_CASE_		:     Dict					=					num_hidden_layers
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					num_attention_heads
										SCREAMING_SNAKE_CASE_		:     Any					=					intermediate_size

										SCREAMING_SNAKE_CASE_		:     List[Any]					=					hidden_dropout_prob
										SCREAMING_SNAKE_CASE_		:     List[str]					=					attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE_		:     Dict					=					max_position_embeddings
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					eos_token_id
										SCREAMING_SNAKE_CASE_		:     str					=					pad_token_id
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					bos_token_id






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					ids_tensor([self.batch_size, self.seq_length - 1]			,						self.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size       )			,						1       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					tf.concat([input_ids, eos_tensor]			,						axis=1       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					ids_tensor([self.batch_size, self.seq_length]			,						self.vocab_size       )

										SCREAMING_SNAKE_CASE_		:     Any					=					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			,						)
										SCREAMING_SNAKE_CASE_		:     Dict					=					prepare_blenderbot_inputs_dict(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )
										return config, inputs_dict






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					TFBlenderbotModel(config=lowerCAmelCase__       ).get_decoder()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					inputs_dict['input_ids']

										SCREAMING_SNAKE_CASE_		:     str					=					input_ids[:1, :]
										SCREAMING_SNAKE_CASE_		:     List[str]					=					inputs_dict['attention_mask'][:1, :]
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					inputs_dict['head_mask']
										SCREAMING_SNAKE_CASE_		:     Dict					=					1

										# first forward pass
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(lowerCAmelCase__			,						attention_mask=lowerCAmelCase__			,						head_mask=lowerCAmelCase__			,						use_cache=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     List[str]					=					outputs.to_tuple()

										# create hypothetical next token and extent to next_input_ids
										SCREAMING_SNAKE_CASE_		:     Dict					=					ids_tensor((self.batch_size, 3)			,						config.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tf.cast(ids_tensor((self.batch_size, 3)			,						2       )			,						tf.inta       )

										# append to next input_ids and
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					tf.concat([input_ids, next_tokens]			,						axis=-1       )
										SCREAMING_SNAKE_CASE_		:     Any					=					tf.concat([attention_mask, next_attn_mask]			,						axis=-1       )

										SCREAMING_SNAKE_CASE_		:     str					=					model(lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )[0]
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					model(lowerCAmelCase__			,						attention_mask=lowerCAmelCase__			,						past_key_values=lowerCAmelCase__       )[0]

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

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

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







def 						a__			(				A__,      A__,      A__,      A__=None,      A__=None,      A__=None,      A__=None,      A__=None,      ):
			if attention_mask is None:
						SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					tf.cast(tf.math.not_equal(A__,      config.pad_token_id				),      tf.inta				)
			if decoder_attention_mask is None:
						SCREAMING_SNAKE_CASE_		:     Optional[int]					=					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:
						SCREAMING_SNAKE_CASE_		:     Tuple					=					tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
			if decoder_head_mask is None:
						SCREAMING_SNAKE_CASE_		:     Dict					=					tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
			if cross_attn_head_mask is None:
						SCREAMING_SNAKE_CASE_		:     Optional[int]					=					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 	__lowercase    (__SCREAMING_SNAKE_CASE	,       __SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							_UpperCAmelCase      =						(TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
							_UpperCAmelCase      =						(TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
							_UpperCAmelCase      =						(
							    {
							        """conversational""": TFBlenderbotForConditionalGeneration,
							        """feature-extraction""": TFBlenderbotModel,
							        """summarization""": TFBlenderbotForConditionalGeneration,
							        """text2text-generation""": TFBlenderbotForConditionalGeneration,
							        """translation""": TFBlenderbotForConditionalGeneration,
							    }
							    if is_tf_available()
							    else {}
							)
							_UpperCAmelCase      =						True
							_UpperCAmelCase      =						False
							_UpperCAmelCase      =						False






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					TFBlenderbotModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					ConfigTester(self			,						config_class=lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										self.config_tester.run_common_tests()






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.model_tester.prepare_config_and_inputs_for_common()
										self.model_tester.check_decoder_model_past_large_inputs(*lowerCAmelCase__       )







@require_tokenizers
@require_tf
class 	__lowercase    (unittest.TestCase  ):





							"""simple docstring"""







							_UpperCAmelCase      =						["""My friends are cool but they eat too many carbs."""]
							_UpperCAmelCase      =						"""facebook/blenderbot-400M-distill"""






							@cached_property
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return BlenderbotTokenizer.from_pretrained(self.model_name       )






							@cached_property
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name       )
										return model






							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.tokenizer(self.src_text			,						return_tensors='tf'       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.model.generate(
										    model_inputs.input_ids			,						)
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.tokenizer.batch_decode(generated_ids.numpy()			,						skip_special_tokens=lowerCAmelCase__       )[0]
										assert (
										    generated_words
										    == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
										)









import itertools
import json
import linecache
import os
import pickle
import re
import socket
import string
from collections import Counter
from logging import getLogger
from pathlib import Path
from typing import Callable, Dict, Iterable, List

import git
import torch
from torch.utils.data import Dataset

from transformers import BartTokenizer, RagTokenizer, TaTokenizer






def 						a__			(				A__,      A__,      A__,      A__,      A__=True,      A__="pt"				):
			SCREAMING_SNAKE_CASE_		:     Tuple					=					{'add_prefix_space': True} if isinstance(A__,      A__				) and not line.startswith(' '				) else {}
			SCREAMING_SNAKE_CASE_		:     Dict					=					padding_side
			return tokenizer(
			    [line],      max_length=A__,      padding='max_length' if pad_to_max_length else None,      truncation=A__,      return_tensors=A__,      add_special_tokens=A__,      **A__,      )






def 						a__			(				A__,      A__,      A__=None,      ):
			SCREAMING_SNAKE_CASE_		:     int					=					input_ids.ne(A__				).any(dim=0				)
			if attention_mask is None:
						return input_ids[:, keep_column_mask]
			else:
						return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__="train"			,						lowerCAmelCase__=None			,						lowerCAmelCase__=None			,						lowerCAmelCase__=None			,						lowerCAmelCase__=""			,						):





										"""simple docstring"""




										super().__init__()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					Path(lowerCAmelCase__       ).joinpath(type_path + '.source'       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					Path(lowerCAmelCase__       ).joinpath(type_path + '.target'       )
										SCREAMING_SNAKE_CASE_		:     int					=					self.get_char_lens(self.src_file       )
										SCREAMING_SNAKE_CASE_		:     int					=					max_source_length
										SCREAMING_SNAKE_CASE_		:     Dict					=					max_target_length
										assert min(self.src_lens       ) > 0, F'''found empty line in {self.src_file}'''
										SCREAMING_SNAKE_CASE_		:     Any					=					tokenizer
										SCREAMING_SNAKE_CASE_		:     List[str]					=					prefix
										if n_obs is not None:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.src_lens[:n_obs]
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					src_lang
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tgt_lang






							def __len__(      self       ):





										"""simple docstring"""




										return len(self.src_lens       )






							def __getitem__(      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					index + 1  # linecache starts at 1
										SCREAMING_SNAKE_CASE_		:     List[str]					=					self.prefix + linecache.getline(str(self.src_file       )			,						lowerCAmelCase__       ).rstrip('\n'       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					linecache.getline(str(self.tgt_file       )			,						lowerCAmelCase__       ).rstrip('\n'       )
										assert source_line, F'''empty source line for index {index}'''
										assert tgt_line, F'''empty tgt line for index {index}'''

										# Need to add eos token manually for T5
										if isinstance(self.tokenizer			,						lowerCAmelCase__       ):
													source_line += self.tokenizer.eos_token
													tgt_line += self.tokenizer.eos_token

										# Pad source and target to the right
										SCREAMING_SNAKE_CASE_		:     Any					=					(
										    self.tokenizer.question_encoder if isinstance(self.tokenizer			,						lowerCAmelCase__       ) else self.tokenizer
										)
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.tokenizer.generator if isinstance(self.tokenizer			,						lowerCAmelCase__       ) else self.tokenizer

										SCREAMING_SNAKE_CASE_		:     Dict					=					encode_line(lowerCAmelCase__			,						lowerCAmelCase__			,						self.max_source_length			,						'right'       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					encode_line(lowerCAmelCase__			,						lowerCAmelCase__			,						self.max_target_length			,						'right'       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					source_inputs['input_ids'].squeeze()
										SCREAMING_SNAKE_CASE_		:     Dict					=					target_inputs['input_ids'].squeeze()
										SCREAMING_SNAKE_CASE_		:     str					=					source_inputs['attention_mask'].squeeze()
										return {
										    "input_ids": source_ids,
										    "attention_mask": src_mask,
										    "decoder_input_ids": target_ids,
										}






							@staticmethod
							def 					UpperCamelCase__     (      lowerCAmelCase__       ):





										"""simple docstring"""




										return [len(lowerCAmelCase__       ) for x in Path(lowerCAmelCase__       ).open().readlines()]






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					torch.stack([x['input_ids'] for x in batch]       )
										SCREAMING_SNAKE_CASE_		:     int					=					torch.stack([x['attention_mask'] for x in batch]       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					torch.stack([x['decoder_input_ids'] for x in batch]       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					(
										    self.tokenizer.generator.pad_token_id
										    if isinstance(self.tokenizer			,						lowerCAmelCase__       )
										    else self.tokenizer.pad_token_id
										)
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					(
										    self.tokenizer.question_encoder.pad_token_id
										    if isinstance(self.tokenizer			,						lowerCAmelCase__       )
										    else self.tokenizer.pad_token_id
										)
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					trim_batch(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					trim_batch(lowerCAmelCase__			,						lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					{
										    'input_ids': source_ids,
										    'attention_mask': source_mask,
										    'decoder_input_ids': y,
										}
										return batch







lowerCAmelCase__			:      Tuple								=getLogger(__name__)






def 						a__			(				A__				):
			return list(itertools.chain.from_iterable(A__				)				)






def 						a__			(				A__				):
			SCREAMING_SNAKE_CASE_		:     Optional[int]					=					get_git_info()
			save_json(A__,      os.path.join(A__,      'git_log.json'				)				)






def 						a__			(				A__,      A__,      A__=4,      **A__				):
			with open(A__,      'w'				) as f:
						json.dump(A__,      A__,      indent=A__,      **A__				)






def 						a__			(				A__				):
			with open(A__				) as f:
						return json.load(A__				)






def 						a__			(				):
			SCREAMING_SNAKE_CASE_		:     List[str]					=					git.Repo(search_parent_directories=A__				)
			SCREAMING_SNAKE_CASE_		:     Any					=					{
			    'repo_id': str(A__				),
			    'repo_sha': str(repo.head.object.hexsha				),
			    'repo_branch': str(repo.active_branch				),
			    'hostname': str(socket.gethostname()				),
			}
			return repo_infos






def 						a__			(				A__,      A__				):
			return list(map(A__,      A__				)				)






def 						a__			(				A__,      A__				):
			with open(A__,      'wb'				) as f:
						return pickle.dump(A__,      A__				)






def 						a__			(				A__				):

			def remove_articles(A__				):
						return re.sub(r'\b(a|an|the)\b',      ' ',      A__				)

			def white_space_fix(A__				):
						return " ".join(text.split()				)

			def remove_punc(A__				):
						SCREAMING_SNAKE_CASE_		:     Optional[int]					=					set(string.punctuation				)
						return "".join(ch for ch in text if ch not in exclude				)

			def lower(A__				):
						return text.lower()

			return white_space_fix(remove_articles(remove_punc(lower(A__				)				)				)				)






def 						a__			(				A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					normalize_answer(A__				).split()
			SCREAMING_SNAKE_CASE_		:     List[str]					=					normalize_answer(A__				).split()
			SCREAMING_SNAKE_CASE_		:     Optional[int]					=					Counter(A__				) & Counter(A__				)
			SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					sum(common.values()				)
			if num_same == 0:
						return 0
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					1.0 * num_same / len(A__				)
			SCREAMING_SNAKE_CASE_		:     int					=					1.0 * num_same / len(A__				)
			SCREAMING_SNAKE_CASE_		:     List[str]					=					(2 * precision * recall) / (precision + recall)
			return fa






def 						a__			(				A__,      A__				):
			return normalize_answer(A__				) == normalize_answer(A__				)






def 						a__			(				A__,      A__				):
			assert len(A__				) == len(A__				)
			SCREAMING_SNAKE_CASE_		:     Any					=					0
			for hypo, pred in zip(A__,      A__				):
						em += exact_match_score(A__,      A__				)
			if len(A__				) > 0:
						em /= len(A__				)
			return {"em": em}






def 						a__			(				A__				):
			return model_prefix.startswith('rag'				)






def 						a__			(				A__,      A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					{p: p for p in extra_params}
			# T5 models don't have `dropout` param, they have `dropout_rate` instead
			SCREAMING_SNAKE_CASE_		:     List[Any]					=					'dropout_rate'
			for p in extra_params:
						if getattr(A__,      A__,      A__				):
									if not hasattr(A__,      A__				) and not hasattr(A__,      equivalent_param[p]				):
												logger.info('config doesn\'t have a `{}` attribute'.format(A__				)				)
												delattr(A__,      A__				)
												continue
									SCREAMING_SNAKE_CASE_		:     Any					=					p if hasattr(A__,      A__				) else equivalent_param[p]
									setattr(A__,      A__,      getattr(A__,      A__				)				)
									delattr(A__,      A__				)
			return hparams, config





