


import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
		    SpeechaTextaDecoder,
		    SpeechaTextaForCausalLM,
		)


@require_torch
class        __lowerCamelCase      :


				"""simple docstring"""







				def __init__(				self     :    Union[str, Any]	,							lowerCamelCase_     :    Optional[Any]	,							lowerCamelCase_     :    List[str]=99	,							lowerCamelCase_     :    Tuple=13	,							lowerCamelCase_     :    Any=16	,							lowerCamelCase_     :    Optional[int]=7	,							lowerCamelCase_     :    Optional[int]=True	,							lowerCamelCase_     :    str=True	,							lowerCamelCase_     :    List[Any]=True	,							lowerCamelCase_     :    Tuple=False	,							lowerCamelCase_     :    List[Any]=True	,							lowerCamelCase_     :    int=2	,							lowerCamelCase_     :    List[Any]=32	,							lowerCamelCase_     :    Union[str, Any]=4	,							lowerCamelCase_     :    str=4	,							lowerCamelCase_     :    List[Any]=30	,							lowerCamelCase_     :    str=0	,							lowerCamelCase_     :    Any=1	,							lowerCamelCase_     :    Optional[Any]=2	,							lowerCamelCase_     :    str=None	,							):
								_lowerCAmelCase  						=parent
								_lowerCAmelCase  						=batch_size
								_lowerCAmelCase  						=decoder_seq_length
								# For common tests
								_lowerCAmelCase  						=self.decoder_seq_length
								_lowerCAmelCase  						=is_training
								_lowerCAmelCase  						=use_attention_mask
								_lowerCAmelCase  						=use_labels

								_lowerCAmelCase  						=vocab_size
								_lowerCAmelCase  						=d_model
								_lowerCAmelCase  						=d_model
								_lowerCAmelCase  						=decoder_layers
								_lowerCAmelCase  						=decoder_layers
								_lowerCAmelCase  						=decoder_ffn_dim
								_lowerCAmelCase  						=decoder_attention_heads
								_lowerCAmelCase  						=decoder_attention_heads
								_lowerCAmelCase  						=eos_token_id
								_lowerCAmelCase  						=bos_token_id
								_lowerCAmelCase  						=pad_token_id
								_lowerCAmelCase  						=decoder_start_token_id
								_lowerCAmelCase  						=use_cache
								_lowerCAmelCase  						=max_position_embeddings

								_lowerCAmelCase  						=None
								_lowerCAmelCase  						=decoder_seq_length
								_lowerCAmelCase  						=2
								_lowerCAmelCase  						=1






				def        lowerCAmelCase__			(				self     :    Optional[Any]					):
								_lowerCAmelCase  						=ids_tensor([self.batch_size, self.decoder_seq_length]	,							self.vocab_size					)

								_lowerCAmelCase  						=None
								if self.use_attention_mask:
												_lowerCAmelCase  						=ids_tensor([self.batch_size, self.decoder_seq_length]	,							vocab_size=2					)

								_lowerCAmelCase  						=None
								if self.use_labels:
												_lowerCAmelCase  						=ids_tensor([self.batch_size, self.decoder_seq_length]	,							self.vocab_size					)

								_lowerCAmelCase  						=SpeechaTextaConfig(
								    vocab_size=self.vocab_size	,							d_model=self.d_model	,							decoder_layers=self.decoder_layers	,							decoder_ffn_dim=self.decoder_ffn_dim	,							decoder_attention_heads=self.decoder_attention_heads	,							eos_token_id=self.eos_token_id	,							bos_token_id=self.bos_token_id	,							use_cache=self.use_cache	,							pad_token_id=self.pad_token_id	,							decoder_start_token_id=self.decoder_start_token_id	,							max_position_embeddings=self.max_position_embeddings	,							)

								return (
								    config,
								    input_ids,
								    attention_mask,
								    lm_labels,
								)






				def        lowerCAmelCase__			(				self     :    str	,							lowerCamelCase_     :    Any	,							lowerCamelCase_     :    List[str]	,							lowerCamelCase_     :    List[Any]	,							lowerCamelCase_     :    Dict	,							):
								_lowerCAmelCase  						=True
								_lowerCAmelCase  						=SpeechaTextaDecoder(config=lowerCamelCase_					).to(lowerCamelCase_					).eval()
								_lowerCAmelCase  						=input_ids[:2]

								input_ids[input_ids == 0] += 1
								# first forward pass
								_lowerCAmelCase  						=model(lowerCamelCase_	,							use_cache=lowerCamelCase_					)
								_lowerCAmelCase  						=model(lowerCamelCase_					)
								_lowerCAmelCase  						=model(lowerCamelCase_	,							use_cache=lowerCamelCase_					)

								self.parent.assertTrue(len(lowerCamelCase_					) == len(lowerCamelCase_					)					)
								self.parent.assertTrue(len(lowerCamelCase_					) == len(lowerCamelCase_					) + 1					)

								_lowerCAmelCase  						=outputs["""past_key_values"""]

								# create hypothetical next token and extent to next_input_ids
								_lowerCAmelCase  						=ids_tensor((2, 1)	,							config.vocab_size - 1					) + 1

								# append to next input_ids and
								_lowerCAmelCase  						=torch.cat([input_ids, next_tokens]	,							dim=-1					)

								print(lowerCamelCase_					)

								_lowerCAmelCase  						=model(lowerCamelCase_					)["""last_hidden_state"""]
								_lowerCAmelCase  						=model(lowerCamelCase_	,							past_key_values=lowerCamelCase_					)["""last_hidden_state"""]

								# select random slice
								_lowerCAmelCase  						=ids_tensor((1,)	,							output_from_past.shape[-1]					).item()
								_lowerCAmelCase  						=output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
								_lowerCAmelCase  						=output_from_past[:, 0, random_slice_idx].detach()

								# test that outputs are equal for slice
								assert torch.allclose(lowerCamelCase_	,							lowerCamelCase_	,							atol=1e-3					)






				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								_lowerCAmelCase  						=self.prepare_config_and_inputs()
								(
								    (
								    _lowerCAmelCase
								)						,							(
								    _lowerCAmelCase
								)						,							(
								    _lowerCAmelCase
								)						,							(
								    _lowerCAmelCase
								)						,							
								)  						=config_and_inputs

								_lowerCAmelCase  						={
								    """input_ids""": input_ids,
								    """attention_mask""": attention_mask,
								}
								return config, inputs_dict


@require_torch
class        __lowerCamelCase      (				lowerCamelCase_				,		lowerCamelCase_				,		lowerCamelCase_				,		unittest.TestCase ):


				"""simple docstring"""







				a_:		List[str]					      =	(SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
				a_:		List[str]					      =	(SpeechaTextaForCausalLM,) if is_torch_available() else ()
				a_:		List[Any]					      =	{"""text-generation""": SpeechaTextaForCausalLM} if is_torch_available() else {}
				a_:		Tuple					      =	True
				a_:		List[Any]					      =	False






				def        lowerCAmelCase__			(				self     :    str	,							):
								_lowerCAmelCase  						=SpeechaTextaStandaloneDecoderModelTester(self	,							is_training=lowerCamelCase_					)
								_lowerCAmelCase  						=ConfigTester(self	,							config_class=lowerCamelCase_					)






				def        lowerCAmelCase__			(				self     :    str					):
								pass






				def        lowerCAmelCase__			(				self     :    Dict					):
								pass






				def        lowerCAmelCase__			(				self     :    str					):
								pass






				def        lowerCAmelCase__			(				self     :    Union[str, Any]					):
								self.config_tester.run_common_tests()






				def        lowerCAmelCase__			(				self     :    str					):
								_lowerCAmelCase  						=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_decoder_model_past(*lowerCamelCase_					)






				def        lowerCAmelCase__			(				self     :    List[str]					):
								return





from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    torch.Tensor       ,     lowercase__     :    bool     ):






				'''simple docstring'''






				if not inplace:
								_lowerCAmelCase  						=ma + ma
				else:
								ma += ma

				return ma

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    List[int]     ):






				'''simple docstring'''






				_lowerCAmelCase  						=-1 * len(lowercase__     )
				_lowerCAmelCase  						=list(range(len(tensor.shape[:zero_index]     )     )     )
				return tensor.permute(first_inds + [zero_index + i for i in inds]     )

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    int     ):






				'''simple docstring'''






				return t.reshape(t.shape[:-no_dims] + (-1,)     )

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    torch.Tensor       ,     lowercase__     :    int       ,     lowercase__     :    float = 1e-4     ):






				'''simple docstring'''






				_lowerCAmelCase  						=mask.expand(*value.shape     )
				return torch.sum(mask * value       ,     dim=lowercase__     ) / (eps + torch.sum(lowercase__       ,     dim=lowercase__     ))

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    torch.types.Number = 2.3_1_2_5       ,     lowercase__     :    torch.types.Number = 2_1.6_8_7_5       ,     lowercase__     :    int = 64     ):






				'''simple docstring'''






				_lowerCAmelCase  						=torch.linspace(lowercase__       ,     lowercase__       ,     no_bins - 1       ,     device=pts.device     )
				_lowerCAmelCase  						=torch.sqrt(torch.sum((pts.unsqueeze(-2     ) - pts.unsqueeze(-3     )) ** 2       ,     dim=-1     )     )
				return torch.bucketize(lowercase__       ,     lowercase__     )

def 		snake_case_   (				lowercase__     :    Callable[[list], Any]       ,     lowercase__     :    List[dict]     ):






				'''simple docstring'''






				_lowerCAmelCase  						=dicts[0]
				_lowerCAmelCase  						={}
				for k, v in first.items():
								_lowerCAmelCase  						=[d[k] for d in dicts]
								if isinstance(lowercase__       ,     lowercase__     ):
												_lowerCAmelCase  						=dict_multimap(lowercase__       ,     lowercase__     )
								else:
												_lowerCAmelCase  						=fn(lowercase__     )

				return new_dict

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    torch.Tensor     ):






				'''simple docstring'''






				_lowerCAmelCase  						=v_bins.view(((1,) * len(x.shape     )) + (len(lowercase__     ),)     )
				_lowerCAmelCase  						=x[..., None] - reshaped_bins
				_lowerCAmelCase  						=torch.argmin(torch.abs(lowercase__     )       ,     dim=-1     )
				return nn.functional.one_hot(lowercase__       ,     num_classes=len(lowercase__     )     ).float()

def 		snake_case_   (				lowercase__     :    torch.Tensor       ,     lowercase__     :    torch.Tensor       ,     lowercase__     :    int = 0       ,     lowercase__     :    int = 0     ):






				'''simple docstring'''






				_lowerCAmelCase  						=[]
				for i, s in enumerate(data.shape[:no_batch_dims]     ):
								_lowerCAmelCase  						=torch.arange(lowercase__     )
								_lowerCAmelCase  						=r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape     ) - i - 1)))     )
								ranges.append(lowercase__     )

				_lowerCAmelCase  						=[slice(lowercase__     ) for _ in range(len(data.shape     ) - no_batch_dims     )]
				_lowerCAmelCase  						=inds
				ranges.extend(lowercase__     )
				# Matt note: Editing this to get around the behaviour of using a list as an array index changing
				# in recent Numpy versions
				return data[tuple(lowercase__     )]


__SCREAMING_SNAKE_CASE	:							str           =  TypeVar('''T''')

def 		snake_case_   (				lowercase__     :    Callable[[T], Any]       ,     lowercase__     :    Dict[Any, Union[dict, list, tuple, T]]       ,     lowercase__     :    Type[T]     ):






				'''simple docstring'''






				_lowerCAmelCase  						={}
				for k, v in dic.items():
								if isinstance(lowercase__       ,     lowercase__     ):
												_lowerCAmelCase  						=dict_map(lowercase__       ,     lowercase__       ,     lowercase__     )
								else:
												_lowerCAmelCase  						=tree_map(lowercase__       ,     lowercase__       ,     lowercase__     )

				return new_dict

@overload
def 		snake_case_   (				lowercase__     :    Callable[[T], Any]       ,     lowercase__     :    T       ,     lowercase__     :    Type[T]     ):






				'''simple docstring'''






				...

@overload
def 		snake_case_   (				lowercase__     :    Callable[[T], Any]       ,     lowercase__     :    dict       ,     lowercase__     :    Type[T]     ):






				'''simple docstring'''






				...

@overload
def 		snake_case_   (				lowercase__     :    Callable[[T], Any]       ,     lowercase__     :    list       ,     lowercase__     :    Type[T]     ):






				'''simple docstring'''






				...

@overload
def 		snake_case_   (				lowercase__     :    Callable[[T], Any]       ,     lowercase__     :    tuple       ,     lowercase__     :    Type[T]     ):






				'''simple docstring'''






				...

def 		snake_case_   (				lowercase__     :    Optional[Any]       ,     lowercase__     :    str       ,     lowercase__     :    Optional[int]     ):






				'''simple docstring'''






				if isinstance(lowercase__       ,     lowercase__     ):
								return dict_map(lowercase__       ,     lowercase__       ,     lowercase__     )
				elif isinstance(lowercase__       ,     lowercase__     ):
								return [tree_map(lowercase__       ,     lowercase__       ,     lowercase__     ) for x in tree]
				elif isinstance(lowercase__       ,     lowercase__     ):
								return tuple(tree_map(lowercase__       ,     lowercase__       ,     lowercase__     ) for x in tree     )
				elif isinstance(lowercase__       ,     lowercase__     ):
								return fn(lowercase__     )
				else:
								print(type(lowercase__     )     )
								raise ValueError("""Not supported"""     )





__SCREAMING_SNAKE_CASE	:							Any           =  partial(tree_map, leaf_type=torch.Tensor)

