
"""simple docstring"""





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, 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
				_a														=		"""platform"""

				import jax
				import jax.numpy as jnp

				from transformers import BlenderbotTokenizer
				from transformers.models.blenderbot.modeling_flax_blenderbot import (
				    FlaxBlenderbotForConditionalGeneration,
				    FlaxBlenderbotModel,
				    shift_tokens_right,
				)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case=None,     __snake_case=None,     __snake_case=None,     __snake_case=None,     __snake_case=None,     __snake_case=None,     )							->		Optional[Any]:


					"""simple docstring"""





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



		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				,		)       ->       int:






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



		def 						UpperCAmelCase							(					self)       ->       str:






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

							_UpperCamelCase							 =  shift_tokens_right(__a				,		1				,		2)

							_UpperCamelCase							 =  BlenderbotConfig(
							    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				,		)
							_UpperCamelCase							 =  prepare_blenderbot_inputs_dict(__a				,		__a				,		__a)
							return config, inputs_dict



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






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self.prepare_config_and_inputs()
							return config, inputs_dict



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  20
							_UpperCamelCase							 =  model_class_name(__a)

							_UpperCamelCase							 =  model.encode(inputs_dict['''input_ids'''])

							_UpperCamelCase ,  _UpperCamelCase							 =  (
							    inputs_dict['''decoder_input_ids'''],
							    inputs_dict['''decoder_attention_mask'''],
							)

							_UpperCamelCase							 =  model.init_cache(decoder_input_ids.shape[0]				,		__a				,		__a)
							_UpperCamelCase							 =  jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,		dtype='''i4''')

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

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

							_UpperCamelCase							 =  model.decode(__a				,		__a)

							_UpperCamelCase							 =  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 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  20
							_UpperCamelCase							 =  model_class_name(__a)

							_UpperCamelCase							 =  model.encode(inputs_dict['''input_ids'''])

							_UpperCamelCase ,  _UpperCamelCase							 =  (
							    inputs_dict['''decoder_input_ids'''],
							    inputs_dict['''decoder_attention_mask'''],
							)

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

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

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

							_UpperCamelCase							 =  model.decode(__a				,		__a				,		decoder_attention_mask=__a)

							_UpperCamelCase							 =  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     _UpperCAmelCase(    unittest.TestCase	):
		lowercase__						=      99



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






							'''simple docstring'''
							_UpperCamelCase							 =  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				,		)

							_UpperCamelCase							 =  input_ids.shape[0]
							_UpperCamelCase							 =  BlenderbotConfig(
							    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 						UpperCAmelCase							(					self)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  self._get_config_and_data()
							_UpperCamelCase							 =  FlaxBlenderbotForConditionalGeneration(__a)
							_UpperCamelCase							 =  lm_model(input_ids=__a)
							_UpperCamelCase							 =  (batch_size, input_ids.shape[1], config.vocab_size)
							self.assertEqual(outputs['''logits'''].shape				,		__a)



		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  BlenderbotConfig(
							    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				,		)
							_UpperCamelCase							 =  FlaxBlenderbotForConditionalGeneration(__a)
							_UpperCamelCase							 =  np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,		dtype=np.intaa)
							_UpperCamelCase							 =  np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,		dtype=np.intaa)
							_UpperCamelCase							 =  lm_model(input_ids=__a				,		decoder_input_ids=__a)
							_UpperCamelCase							 =  (*summary.shape, config.vocab_size)
							self.assertEqual(outputs['''logits'''].shape				,		__a)




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






							'''simple docstring'''
							_UpperCamelCase							 =  np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,		dtype=np.intaa)
							_UpperCamelCase							 =  shift_tokens_right(__a				,		1				,		2)
							_UpperCamelCase							 =  np.equal(__a				,		1).astype(np.floataa).sum()
							_UpperCamelCase							 =  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     _UpperCAmelCase(    lowerCamelCase      ,     unittest.TestCase      ,     lowerCamelCase	):
		lowercase__						=      True
		lowercase__						=      (
		    (
		        FlaxBlenderbotModel,
		        FlaxBlenderbotForConditionalGeneration,
		    )
		    if is_flax_available()
		    else ()
		)
		lowercase__						=      (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()



		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  FlaxBlenderbotModelTester(self)



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  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 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  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 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__):
																	_UpperCamelCase							 =  self._prepare_for_class(__a				,		__a)
																	_UpperCamelCase							 =  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'''):
																						_UpperCamelCase							 =  encode_jitted(**__a).to_tuple()

																	with self.subTest('''JIT Disabled'''):
																						with jax.disable_jit():
																											_UpperCamelCase							 =  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 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__):
																	_UpperCamelCase							 =  model_class(__a)
																	_UpperCamelCase							 =  model.encode(inputs_dict['''input_ids''']				,		inputs_dict['''attention_mask'''])

																	_UpperCamelCase							 =  {
																	    '''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'''):
																						_UpperCamelCase							 =  decode_jitted(**__a).to_tuple()

																	with self.subTest('''JIT Disabled'''):
																						with jax.disable_jit():
																											_UpperCamelCase							 =  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 						UpperCAmelCase							(					self)       ->       Union[str, Any]:






							'''simple docstring'''
							for model_class_name in self.all_model_classes:
												_UpperCamelCase							 =  model_class_name.from_pretrained('''facebook/blenderbot-400M-distill''')
												# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
												_UpperCamelCase							 =  np.ones((1, 1)) * model.config.eos_token_id
												_UpperCamelCase							 =  model(__a)
												self.assertIsNotNone(__a)




		@unittest.skipUnless(jax_device != '''cpu'''				,		'''3B test too slow on CPU.''')
		@slow
		def 						UpperCAmelCase							(					self)       ->       Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  {'''num_beams''': 1, '''early_stopping''': True, '''min_length''': 15, '''max_length''': 25}
							_UpperCamelCase							 =  {'''skip_special_tokens''': True, '''clean_up_tokenization_spaces''': True}

							_UpperCamelCase							 =  FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''				,		from_pt=__a)
							_UpperCamelCase							 =  BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B''')

							_UpperCamelCase							 =  ['''Sam''']
							_UpperCamelCase							 =  tokenizer(__a				,		return_tensors='''jax''')

							_UpperCamelCase							 =  model.generate(**__a				,		**__a)
							_UpperCamelCase							 =  '''Sam is a great name. It means "sun" in Gaelic.'''

							_UpperCamelCase							 =  tokenizer.batch_decode(__a				,		**__a)
							assert generated_txt[0].strip() == tgt_text








"""simple docstring"""





from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


_a														=		logging.get_logger(__name__)






class     _UpperCAmelCase(    lowerCamelCase	):



		def 						UpperCAmelCase							(					self				,		__a)       ->       Tuple:






							'''simple docstring'''
							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [label.strip() for label in labels.split(''',''') if label.strip()]
							return labels




		def __call__(					self				,		__a				,		__a				,		__a)       ->       Optional[int]:






							'''simple docstring'''
							if len(__a) == 0 or len(__a) == 0:
												raise ValueError('''You must include at least one label and at least one sequence.''')
							if hypothesis_template.format(labels[0]) == hypothesis_template:
												raise ValueError(
												    (
												        '''The provided hypothesis_template "{}" was not able to be formatted with the target labels. '''
												        '''Make sure the passed template includes formatting syntax such as {{}} where the label should go.'''
												    ).format(__a))

							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [sequences]

							_UpperCamelCase							 =  []
							for sequence in sequences:
												sequence_pairs.extend([[sequence, hypothesis_template.format(__a)] for label in labels])

							return sequence_pairs, sequences






@add_end_docstrings(lowerCamelCase	)
class     _UpperCAmelCase(    lowerCamelCase	):



		def __init__(					self				,		__a=ZeroShotClassificationArgumentHandler()				,		*__a				,		**__a)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  args_parser
							super().__init__(*__a				,		**__a)
							if self.entailment_id == -1:
												logger.warning(
												    '''Failed to determine \'entailment\' label id from the label2id mapping in the model config. Setting to '''
												    '''-1. Define a descriptive label2id mapping in the model config to ensure correct outputs.''')



		@property
		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							for label, ind in self.model.config.labelaid.items():
												if label.lower().startswith('''entail'''):
																	return ind
							return -1



		def 						UpperCAmelCase							(					self				,		__a				,		__a=True				,		__a=True				,		__a=TruncationStrategy.ONLY_FIRST				,		**__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  self.framework
							if self.tokenizer.pad_token is None:
												# Override for tokenizers not supporting padding
												logger.error(
												    '''Tokenizer was not supporting padding necessary for zero-shot, attempting to use '''
												    ''' `pad_token=eos_token`''')
												_UpperCamelCase							 =  self.tokenizer.eos_token
							try:
												_UpperCamelCase							 =  self.tokenizer(
												    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=__a				,		)
							except Exception as e:
												if "too short" in str(__a):
																	# tokenizers might yell that we want to truncate
																	# to a value that is not even reached by the input.
																	# In that case we don't want to truncate.
																	# It seems there's not a really better way to catch that
																	# exception.

																	_UpperCamelCase							 =  self.tokenizer(
																	    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=TruncationStrategy.DO_NOT_TRUNCATE				,		)
												else:
																	raise e

							return inputs



		def 						UpperCAmelCase							(					self				,		**__a)       ->       Any:






							'''simple docstring'''
							if kwargs.get('''multi_class'''				,		__a) is not None:
												_UpperCamelCase							 =  kwargs['''multi_class''']
												logger.warning(
												    '''The `multi_class` argument has been deprecated and renamed to `multi_label`. '''
												    '''`multi_class` will be removed in a future version of Transformers.''')
							_UpperCamelCase							 =  {}
							if "candidate_labels" in kwargs:
												_UpperCamelCase							 =  self._args_parser._parse_labels(kwargs['''candidate_labels'''])
							if "hypothesis_template" in kwargs:
												_UpperCamelCase							 =  kwargs['''hypothesis_template''']

							_UpperCamelCase							 =  {}
							if "multi_label" in kwargs:
												_UpperCamelCase							 =  kwargs['''multi_label''']
							return preprocess_params, {}, postprocess_params



		def __call__(					self				,		__a				,		*__a				,		**__a				,		)       ->       int:






							'''simple docstring'''
							if len(__a) == 0:
												pass
							elif len(__a) == 1 and "candidate_labels" not in kwargs:
												_UpperCamelCase							 =  args[0]
							else:
												raise ValueError(F'''Unable to understand extra arguments {args}''')

							return super().__call__(__a				,		**__a)



		def 						UpperCAmelCase							(					self				,		__a				,		__a=None				,		__a="This example is {}.")       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self._args_parser(__a				,		__a				,		__a)

							for i, (candidate_label, sequence_pair) in enumerate(zip(__a				,		__a)):
												_UpperCamelCase							 =  self._parse_and_tokenize([sequence_pair])

												yield {
												    "candidate_label": candidate_label,
												    "sequence": sequences[0],
												    "is_last": i == len(__a) - 1,
												    **model_input,
												}



		def 						UpperCAmelCase							(					self				,		__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  inputs['''candidate_label''']
							_UpperCamelCase							 =  inputs['''sequence''']
							_UpperCamelCase							 =  {k: inputs[k] for k in self.tokenizer.model_input_names}
							_UpperCamelCase							 =  self.model(**__a)

							_UpperCamelCase							 =  {
							    '''candidate_label''': candidate_label,
							    '''sequence''': sequence,
							    '''is_last''': inputs['''is_last'''],
							    **outputs,
							}
							return model_outputs




		def 						UpperCAmelCase							(					self				,		__a				,		__a=False)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  [outputs['''candidate_label'''] for outputs in model_outputs]
							_UpperCamelCase							 =  [outputs['''sequence'''] for outputs in model_outputs]
							_UpperCamelCase							 =  np.concatenate([output['''logits'''].numpy() for output in model_outputs])
							_UpperCamelCase							 =  logits.shape[0]
							_UpperCamelCase							 =  len(__a)
							_UpperCamelCase							 =  N // n
							_UpperCamelCase							 =  logits.reshape((num_sequences, n, -1))

							if multi_label or len(__a) == 1:
												# softmax over the entailment vs. contradiction dim for each label independently
												_UpperCamelCase							 =  self.entailment_id
												_UpperCamelCase							 =  -1 if entailment_id == 0 else 0
												_UpperCamelCase							 =  reshaped_outputs[..., [contradiction_id, entailment_id]]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)
												_UpperCamelCase							 =  scores[..., 1]
							else:
												# softmax the "entailment" logits over all candidate labels
												_UpperCamelCase							 =  reshaped_outputs[..., self.entailment_id]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)

							_UpperCamelCase							 =  list(reversed(scores[0].argsort()))
							return {
							    "sequence": sequences[0],
							    "labels": [candidate_labels[i] for i in top_inds],
							    "scores": scores[0, top_inds].tolist(),
							}






