




"""simple docstring"""




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, 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
					lowerCAmelCase							:    Dict  							=							"""platform"""

					import jax
					import jax.numpy as jnp

					from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
					    FlaxBlenderbotSmallForConditionalGeneration,
					    FlaxBlenderbotSmallModel,
					    shift_tokens_right,
					)



def 					a__				(       snake_case__ ,				snake_case__ ,				snake_case__=None ,				snake_case__=None ,				snake_case__=None ,				snake_case__=None ,				snake_case__=None ,				snake_case__=None ,				)		->      Tuple:
			if attention_mask is None:
						lowerCamelCase       =      np.where(input_ids != config.pad_token_id ,				1 ,				0       )
			if decoder_attention_mask is None:
						lowerCamelCase       =      np.where(decoder_input_ids != config.pad_token_id ,				1 ,				0       )
			if head_mask is None:
						lowerCamelCase       =      np.ones((config.encoder_layers, config.encoder_attention_heads)       )
			if decoder_head_mask is None:
						lowerCamelCase       =      np.ones((config.decoder_layers, config.decoder_attention_heads)       )
			if cross_attn_head_mask is None:
						lowerCamelCase       =      np.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": attention_mask,
			}






class 					__magic_name__							:






				'''simple docstring'''





				def __init__( self	,				_a	,				_a=13	,				_a=7	,				_a=True	,				_a=False	,				_a=99	,				_a=16	,				_a=2	,				_a=4	,				_a=4	,				_a="gelu"	,				_a=0.1	,				_a=0.1	,				_a=32	,				_a=2	,				_a=1	,				_a=0	,				_a=0.02	,				):





							"""simple docstring"""
							lowerCamelCase       =      parent
							lowerCamelCase       =      batch_size
							lowerCamelCase       =      seq_length
							lowerCamelCase       =      is_training
							lowerCamelCase       =      use_labels
							lowerCamelCase       =      vocab_size
							lowerCamelCase       =      hidden_size
							lowerCamelCase       =      num_hidden_layers
							lowerCamelCase       =      num_attention_heads
							lowerCamelCase       =      intermediate_size
							lowerCamelCase       =      hidden_act
							lowerCamelCase       =      hidden_dropout_prob
							lowerCamelCase       =      attention_probs_dropout_prob
							lowerCamelCase       =      max_position_embeddings
							lowerCamelCase       =      eos_token_id
							lowerCamelCase       =      pad_token_id
							lowerCamelCase       =      bos_token_id
							lowerCamelCase       =      initializer_range





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      np.clip(ids_tensor([self.batch_size, self.seq_length - 1]	,				self.vocab_size				)	,				3	,				self.vocab_size				)
							lowerCamelCase       =      np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)	,				dtype=np.intaa				))	,				-1				)

							lowerCamelCase       =      shift_tokens_right(_a	,				1	,				2				)

							lowerCamelCase       =      BlenderbotSmallConfig(
							    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_id=self.eos_token_id	,				bos_token_id=self.bos_token_id	,				pad_token_id=self.pad_token_id	,				initializer_range=self.initializer_range	,				use_cache=_a	,				)
							lowerCamelCase       =      prepare_blenderbot_inputs_dict(_a	,				_a	,				_a				)
							return config, inputs_dict





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      self.prepare_config_and_inputs()
							return config, inputs_dict





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      20
							lowerCamelCase       =      model_class_name(_a				)

							lowerCamelCase       =      model.encode(inputs_dict["""input_ids"""]				)

							lowerCamelCase   ,     lowerCamelCase       =      (
							    inputs_dict["""decoder_input_ids"""],
							    inputs_dict["""decoder_attention_mask"""],
							)

							lowerCamelCase       =      model.init_cache(decoder_input_ids.shape[0]	,				_a	,				_a				)
							lowerCamelCase       =      jnp.ones((decoder_input_ids.shape[0], max_decoder_length)	,				dtype="""i4"""				)

							lowerCamelCase       =      jnp.broadcast_to(
							    jnp.arange(decoder_input_ids.shape[-1] - 1				)[None, :]	,				(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,				)
							lowerCamelCase       =      model.decode(
							    decoder_input_ids[:, :-1]	,				_a	,				decoder_attention_mask=_a	,				past_key_values=_a	,				decoder_position_ids=_a	,				)

							lowerCamelCase       =      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,				dtype="""i4"""				)
							lowerCamelCase       =      model.decode(
							    decoder_input_ids[:, -1:]	,				_a	,				decoder_attention_mask=_a	,				past_key_values=outputs_cache.past_key_values	,				decoder_position_ids=_a	,				)

							lowerCamelCase       =      model.decode(_a	,				_a				)

							lowerCamelCase       =      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	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      20
							lowerCamelCase       =      model_class_name(_a				)

							lowerCamelCase       =      model.encode(inputs_dict["""input_ids"""]				)

							lowerCamelCase   ,     lowerCamelCase       =      (
							    inputs_dict["""decoder_input_ids"""],
							    inputs_dict["""decoder_attention_mask"""],
							)

							lowerCamelCase       =      jnp.concatenate(
							    [
							        decoder_attention_mask,
							        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])				),
							    ]	,				axis=-1	,				)

							lowerCamelCase       =      model.init_cache(decoder_input_ids.shape[0]	,				_a	,				_a				)
							lowerCamelCase       =      jnp.broadcast_to(
							    jnp.arange(decoder_input_ids.shape[-1] - 1				)[None, :]	,				(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,				)

							lowerCamelCase       =      model.decode(
							    decoder_input_ids[:, :-1]	,				_a	,				decoder_attention_mask=_a	,				past_key_values=_a	,				decoder_position_ids=_a	,				)
							lowerCamelCase       =      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,				dtype="""i4"""				)
							lowerCamelCase       =      model.decode(
							    decoder_input_ids[:, -1:]	,				_a	,				past_key_values=outputs_cache.past_key_values	,				decoder_attention_mask=_a	,				decoder_position_ids=_a	,				)

							lowerCamelCase       =      model.decode(_a	,				_a	,				decoder_attention_mask=_a				)

							lowerCamelCase       =      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 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''
				__UpperCamelCase    						=		99





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      np.array(
							    [
							        [71, 82, 18, 33, 46, 91, 2],
							        [68, 34, 26, 58, 30, 82, 2],
							        [5, 97, 17, 39, 94, 40, 2],
							        [76, 83, 94, 25, 70, 78, 2],
							        [87, 59, 41, 35, 48, 66, 2],
							        [55, 13, 16, 58, 5, 2, 1],  # note padding
							        [64, 27, 31, 51, 12, 75, 2],
							        [52, 64, 86, 17, 83, 39, 2],
							        [48, 61, 9, 24, 71, 82, 2],
							        [26, 1, 60, 48, 22, 13, 2],
							        [21, 5, 62, 28, 14, 76, 2],
							        [45, 98, 37, 86, 59, 48, 2],
							        [70, 70, 50, 9, 28, 0, 2],
							    ]	,				dtype=np.intaa	,				)

							lowerCamelCase       =      input_ids.shape[0]
							lowerCamelCase       =      BlenderbotSmallConfig(
							    vocab_size=self.vocab_size	,				d_model=24	,				encoder_layers=2	,				decoder_layers=2	,				encoder_attention_heads=2	,				decoder_attention_heads=2	,				encoder_ffn_dim=32	,				decoder_ffn_dim=32	,				max_position_embeddings=48	,				eos_token_id=2	,				pad_token_id=1	,				bos_token_id=0	,				)
							return config, input_ids, batch_size





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self._get_config_and_data()
							lowerCamelCase       =      FlaxBlenderbotSmallForConditionalGeneration(_a				)
							lowerCamelCase       =      lm_model(input_ids=_a				)
							lowerCamelCase       =      (batch_size, input_ids.shape[1], config.vocab_size)
							self.assertEqual(outputs["""logits"""].shape	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      BlenderbotSmallConfig(
							    vocab_size=self.vocab_size	,				d_model=14	,				encoder_layers=2	,				decoder_layers=2	,				encoder_attention_heads=2	,				decoder_attention_heads=2	,				encoder_ffn_dim=8	,				decoder_ffn_dim=8	,				max_position_embeddings=48	,				)
							lowerCamelCase       =      FlaxBlenderbotSmallForConditionalGeneration(_a				)
							lowerCamelCase       =      np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]	,				dtype=np.intaa				)
							lowerCamelCase       =      np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]	,				dtype=np.intaa				)
							lowerCamelCase       =      lm_model(input_ids=_a	,				decoder_input_ids=_a				)
							lowerCamelCase       =      (*summary.shape, config.vocab_size)
							self.assertEqual(outputs["""logits"""].shape	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]	,				dtype=np.intaa				)
							lowerCamelCase       =      shift_tokens_right(_a	,				1	,				2				)
							lowerCamelCase       =      np.equal(_a	,				1				).astype(np.floataa				).sum()
							lowerCamelCase       =      np.equal(_a	,				1				).astype(np.floataa				).sum()
							self.assertEqual(shifted.shape	,				input_ids.shape				)
							self.assertEqual(_a	,				n_pad_before - 1				)
							self.assertTrue(np.equal(shifted[:, 0]	,				2				).all()				)






@require_flax
class 					__magic_name__							( UpperCAmelCase__		,      unittest.TestCase		,      UpperCAmelCase__       ):






				'''simple docstring'''
				__UpperCamelCase    						=		True
				__UpperCamelCase    						=		(
				    (
				        FlaxBlenderbotSmallModel,
				        FlaxBlenderbotSmallForConditionalGeneration,
				    )
				    if is_flax_available()
				    else ()
				)
				__UpperCamelCase    						=		(FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      FlaxBlenderbotSmallModelTester(self				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							for model_class in self.all_model_classes:
										self.model_tester.check_use_cache_forward(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      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(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__				):
													lowerCamelCase       =      self._prepare_for_class(_a	,				_a				)
													lowerCamelCase       =      model_class(_a				)

													@jax.jit
													def encode_jitted(_a	,				_a=None	,				**_a				):
																return model.encode(input_ids=_a	,				attention_mask=_a				)

													with self.subTest("""JIT Enabled"""				):
																lowerCamelCase       =      encode_jitted(**_a				).to_tuple()

													with self.subTest("""JIT Disabled"""				):
																with jax.disable_jit():
																			lowerCamelCase       =      encode_jitted(**_a				).to_tuple()

													self.assertEqual(len(_a				)	,				len(_a				)				)
													for jitted_output, output in zip(_a	,				_a				):
																self.assertEqual(jitted_output.shape	,				output.shape				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__				):
													lowerCamelCase       =      model_class(_a				)
													lowerCamelCase       =      model.encode(inputs_dict["""input_ids"""]	,				inputs_dict["""attention_mask"""]				)

													lowerCamelCase       =      {
													    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
													    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
													    """encoder_outputs""": encoder_outputs,
													}

													@jax.jit
													def decode_jitted(_a	,				_a	,				_a				):
																return model.decode(
																    decoder_input_ids=_a	,				decoder_attention_mask=_a	,				encoder_outputs=_a	,				)

													with self.subTest("""JIT Enabled"""				):
																lowerCamelCase       =      decode_jitted(**_a				).to_tuple()

													with self.subTest("""JIT Disabled"""				):
																with jax.disable_jit():
																			lowerCamelCase       =      decode_jitted(**_a				).to_tuple()

													self.assertEqual(len(_a				)	,				len(_a				)				)
													for jitted_output, output in zip(_a	,				_a				):
																self.assertEqual(jitted_output.shape	,				output.shape				)





				@slow
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							for model_class_name in self.all_model_classes:
										lowerCamelCase       =      model_class_name.from_pretrained("""facebook/blenderbot_small-90M"""				)
										# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
										lowerCamelCase       =      np.ones((1, 1)				) * model.config.eos_token_id
										lowerCamelCase       =      model(_a				)
										self.assertIsNotNone(_a				)









"""simple docstring"""




import uuid
from typing import Any, Dict, List, Optional, Union

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch


lowerCAmelCase							:    List[str]  							=							logging.get_logger(__name__)






class 					__magic_name__							:






				'''simple docstring'''





				def __init__( self	,				_a = None	,				_a = None	,				_a=None	,				_a=None				):





							"""simple docstring"""
							if not conversation_id:
										lowerCamelCase       =      uuid.uuida()
							if past_user_inputs is None:
										lowerCamelCase       =      []
							if generated_responses is None:
										lowerCamelCase       =      []

							lowerCamelCase       =      conversation_id
							lowerCamelCase       =      past_user_inputs
							lowerCamelCase       =      generated_responses
							lowerCamelCase       =      text





				def __eq__( self	,				_a				):





							"""simple docstring"""
							if not isinstance(_a	,				_a				):
										return False
							if self.uuid == other.uuid:
										return True
							return (
							    self.new_user_input == other.new_user_input
							    and self.past_user_inputs == other.past_user_inputs
							    and self.generated_responses == other.generated_responses
							)





				def      _lowerCAmelCase       ( self	,				_a	,				_a = False				):





							"""simple docstring"""
							if self.new_user_input:
										if overwrite:
													logger.warning(
													    f'User input added while unprocessed input was existing: "{self.new_user_input}" was overwritten '
													    f'with: "{text}".'				)
													lowerCamelCase       =      text
										else:
													logger.warning(
													    f'User input added while unprocessed input was existing: "{self.new_user_input}" new input '
													    f'ignored: "{text}". Set `overwrite` to True to overwrite unprocessed user input'				)
							else:
										lowerCamelCase       =      text





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							if self.new_user_input:
										self.past_user_inputs.append(self.new_user_input				)
							lowerCamelCase       =      None





				def      _lowerCAmelCase       ( self	,				_a				):





							"""simple docstring"""
							self.generated_responses.append(_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							for user_input, generated_response in zip(self.past_user_inputs	,				self.generated_responses				):
										yield True, user_input
										yield False, generated_response
							if self.new_user_input:
										yield True, self.new_user_input





				def __repr__( self				):





							"""simple docstring"""
							lowerCamelCase       =      f'Conversation id: {self.uuid} \n'
							for is_user, text in self.iter_texts():
										lowerCamelCase       =      """user""" if is_user else """bot"""
										output += f'{name} >> {text} \n'
							return output






@add_end_docstrings(
    UpperCAmelCase__		,      R"\n        min_length_for_response (`int`, *optional*, defaults to 32):\n            The minimum length (in number of tokens) for a response.\n        minimum_tokens (`int`, *optional*, defaults to 10):\n            The minimum length of tokens to leave for a response.\n    "		,      )
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def __init__( self	,				*_a	,				**_a				):





							"""simple docstring"""
							super().__init__(*_a	,				**_a				)
							if self.tokenizer.pad_token_id is None:
										lowerCamelCase       =      self.tokenizer.eos_token





				def      _lowerCAmelCase       ( self	,				_a=None	,				_a=None	,				_a=None	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      {}
							lowerCamelCase       =      {}
							lowerCamelCase       =      {}

							if min_length_for_response is not None:
										lowerCamelCase       =      min_length_for_response
							if minimum_tokens is not None:
										lowerCamelCase       =      minimum_tokens

							if "max_length" in generate_kwargs:
										lowerCamelCase       =      generate_kwargs["""max_length"""]
										# self.max_length = generate_kwargs.get("max_length", self.model.config.max_length)
							if clean_up_tokenization_spaces is not None:
										lowerCamelCase       =      clean_up_tokenization_spaces

							if generate_kwargs:
										forward_params.update(_a				)
							return preprocess_params, forward_params, postprocess_params





				def __call__( self	,				_a	,				_a=0	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      super().__call__(_a	,				num_workers=_a	,				**_a				)
							if isinstance(_a	,				_a				) and len(_a				) == 1:
										return outputs[0]
							return outputs





				def      _lowerCAmelCase       ( self	,				_a	,				_a=32				):





							"""simple docstring"""
							if not isinstance(_a	,				_a				):
										raise ValueError("""ConversationalPipeline, expects Conversation as inputs"""				)
							if conversation.new_user_input is None:
										raise ValueError(
										    f'Conversation with UUID {type(conversation.uuid				)} does not contain new user input to process. '
										    """Add user inputs with the conversation's `add_user_input` method"""				)
							if hasattr(self.tokenizer	,				"""_build_conversation_input_ids"""				):
										lowerCamelCase       =      self.tokenizer._build_conversation_input_ids(_a				)
							else:
										# If the tokenizer cannot handle conversations, we default to only the old version
										lowerCamelCase       =      self._legacy_parse_and_tokenize(_a				)

							if self.framework == "pt":
										lowerCamelCase       =      torch.LongTensor([input_ids]				)
							elif self.framework == "tf":
										lowerCamelCase       =      tf.constant([input_ids]				)
							return {"input_ids": input_ids, "conversation": conversation}





				def      _lowerCAmelCase       ( self	,				_a	,				_a=10	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      generate_kwargs.get("""max_length"""	,				self.model.config.max_length				)

							lowerCamelCase       =      model_inputs["""input_ids"""].shape[1]
							if max_length - minimum_tokens < n:
										logger.warning(f'Conversation input is to long ({n}), trimming it to ({max_length} - {minimum_tokens})'				)
										lowerCamelCase       =      max_length - minimum_tokens
										lowerCamelCase       =      model_inputs["""input_ids"""][:, -trim:]
										if "attention_mask" in model_inputs:
													lowerCamelCase       =      model_inputs["""attention_mask"""][:, -trim:]
							lowerCamelCase       =      model_inputs.pop("""conversation"""				)
							lowerCamelCase       =      max_length
							lowerCamelCase       =      self.model.generate(**_a	,				**_a				)
							if self.model.config.is_encoder_decoder:
										lowerCamelCase       =      1
							else:
										lowerCamelCase       =      n
							return {"output_ids": output_ids[:, start_position:], "conversation": conversation}





				def      _lowerCAmelCase       ( self	,				_a	,				_a=True				):





							"""simple docstring"""
							lowerCamelCase       =      model_outputs["""output_ids"""]
							lowerCamelCase       =      self.tokenizer.decode(
							    output_ids[0]	,				skip_special_tokens=_a	,				clean_up_tokenization_spaces=_a	,				)
							lowerCamelCase       =      model_outputs["""conversation"""]
							conversation.mark_processed()
							conversation.append_response(_a				)
							return conversation





				def      _lowerCAmelCase       ( self	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.tokenizer.eos_token_id
							lowerCamelCase       =      []
							for is_user, text in conversation.iter_texts():
										if eos_token_id is not None:
													input_ids.extend(self.tokenizer.encode(_a	,				add_special_tokens=_a				) + [eos_token_id]				)
										else:
													input_ids.extend(self.tokenizer.encode(_a	,				add_special_tokens=_a				)				)

							if len(_a				) > self.tokenizer.model_max_length:
										lowerCamelCase       =      input_ids[-self.tokenizer.model_max_length :]
							return input_ids



