
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    _UpperCamelCase:

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

													_UpperCAmelCase  :			int	       =		vocab_size
													_UpperCAmelCase  :			List[Any]	       =		d_model
													_UpperCAmelCase  :			Any	       =		d_model
													_UpperCAmelCase  :			Dict	       =		decoder_layers
													_UpperCAmelCase  :			Optional[Any]	       =		decoder_layers
													_UpperCAmelCase  :			Optional[Any]	       =		decoder_ffn_dim
													_UpperCAmelCase  :			Optional[Any]	       =		decoder_attention_heads
													_UpperCAmelCase  :			Union[str, Any]	       =		decoder_attention_heads
													_UpperCAmelCase  :			int	       =		eos_token_id
													_UpperCAmelCase  :			List[Any]	       =		bos_token_id
													_UpperCAmelCase  :			List[Any]	       =		pad_token_id
													_UpperCAmelCase  :			Optional[int]	       =		decoder_start_token_id
													_UpperCAmelCase  :			List[str]	       =		use_cache
													_UpperCAmelCase  :			int	       =		max_position_embeddings

													_UpperCAmelCase  :			List[str]	       =		None
													_UpperCAmelCase  :			List[str]	       =		decoder_seq_length
													_UpperCAmelCase  :			Union[str, Any]	       =		2
													_UpperCAmelCase  :			Optional[int]	       =		1

						def 						a__ (					self		:			Tuple			):
													_UpperCAmelCase  :			List[str]	       =		ids_tensor([self.batch_size, self.decoder_seq_length]	,						self.vocab_size			)

													_UpperCAmelCase  :			Union[str, Any]	       =		None
													if self.use_attention_mask:
																				_UpperCAmelCase  :			str	       =		ids_tensor([self.batch_size, self.decoder_seq_length]	,						vocab_size=2			)

													_UpperCAmelCase  :			Dict	       =		None
													if self.use_labels:
																				_UpperCAmelCase  :			Union[str, Any]	       =		ids_tensor([self.batch_size, self.decoder_seq_length]	,						self.vocab_size			)

													_UpperCAmelCase  :			Tuple	       =		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 						a__ (					self		:			Tuple	,						_lowerCamelCase		:			Optional[int]	,						_lowerCamelCase		:			List[Any]	,						_lowerCamelCase		:			Any	,						_lowerCamelCase		:			Dict	,						):
													_UpperCAmelCase  :			str	       =		True
													_UpperCAmelCase  :			int	       =		SpeechaTextaDecoder(config=_UpperCAmelCase			).to(_UpperCAmelCase			).eval()
													_UpperCAmelCase  :			List[str]	       =		input_ids[:2]

													input_ids[input_ids == 0] += 1
													# first forward pass
													_UpperCAmelCase  :			List[Any]	       =		model(_UpperCAmelCase	,						use_cache=_UpperCAmelCase			)
													_UpperCAmelCase  :			Union[str, Any]	       =		model(_UpperCAmelCase			)
													_UpperCAmelCase  :			int	       =		model(_UpperCAmelCase	,						use_cache=_UpperCAmelCase			)

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

													_UpperCAmelCase  :			Tuple	       =		outputs["past_key_values"]

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

													# append to next input_ids and
													_UpperCAmelCase  :			Union[str, Any]	       =		torch.cat([input_ids, next_tokens]	,						dim=-1			)

													print(_UpperCAmelCase			)

													_UpperCAmelCase  :			str	       =		model(_UpperCAmelCase			)["last_hidden_state"]
													_UpperCAmelCase  :			List[str]	       =		model(_UpperCAmelCase	,						past_key_values=_UpperCAmelCase			)["last_hidden_state"]

													# select random slice
													_UpperCAmelCase  :			List[str]	       =		ids_tensor((1,)	,						output_from_past.shape[-1]			).item()
													_UpperCAmelCase  :			Optional[Any]	       =		output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
													_UpperCAmelCase  :			int	       =		output_from_past[:, 0, random_slice_idx].detach()

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






						def 						a__ (					self		:			str			):
													_UpperCAmelCase  :			Tuple	       =		self.prepare_config_and_inputs()
													(
													    (
													    _UpperCAmelCase   
													)       ,(
													    _UpperCAmelCase   
													)       ,(
													    _UpperCAmelCase   
													)       ,(
													    _UpperCAmelCase   
													)       ,   
													)  :			Dict	       =		config_and_inputs

													_UpperCAmelCase  :			Optional[Any]	       =		{
													    "input_ids": input_ids,
													    "attention_mask": attention_mask,
													}
													return config, inputs_dict







@require_torch
class    _UpperCamelCase(   lowerCAmelCase__		,						lowerCAmelCase__		,						lowerCAmelCase__		,						unittest.TestCase				):
						__A:						Tuple				   =      (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
						__A:						Optional[int]				   =      (SpeechaTextaForCausalLM,) if is_torch_available() else ()
						__A:						str				   =      {"text-generation": SpeechaTextaForCausalLM} if is_torch_available() else {}
						__A:						Union[str, Any]				   =      True
						__A:						Dict				   =      False

						def 						a__ (					self		:			Dict	,						):
													_UpperCAmelCase  :			int	       =		SpeechaTextaStandaloneDecoderModelTester(self	,						is_training=_UpperCAmelCase			)
													_UpperCAmelCase  :			Dict	       =		ConfigTester(self	,						config_class=_UpperCAmelCase			)

						def 						a__ (					self		:			Dict			):
													pass

						def 						a__ (					self		:			Optional[int]			):
													pass

						def 						a__ (					self		:			Dict			):
													pass

						def 						a__ (					self		:			Any			):
													self.config_tester.run_common_tests()

						def 						a__ (					self		:			Tuple			):
													_UpperCAmelCase  :			Optional[int]	       =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_decoder_model_past(*_UpperCAmelCase			)






						def 						a__ (					self		:			Tuple			):
													return






# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ..utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


SCREAMING_SNAKE_CASE__				=							{
    """agents""": ["""Agent""", """AzureOpenAiAgent""", """HfAgent""", """LocalAgent""", """OpenAiAgent"""],
    """base""": ["""PipelineTool""", """RemoteTool""", """Tool""", """launch_gradio_demo""", """load_tool"""],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__				=							["""DocumentQuestionAnsweringTool"""]
							SCREAMING_SNAKE_CASE__				=							["""ImageCaptioningTool"""]
							SCREAMING_SNAKE_CASE__				=							["""ImageQuestionAnsweringTool"""]
							SCREAMING_SNAKE_CASE__				=							["""ImageSegmentationTool"""]
							SCREAMING_SNAKE_CASE__				=							["""SpeechToTextTool"""]
							SCREAMING_SNAKE_CASE__				=							["""TextClassificationTool"""]
							SCREAMING_SNAKE_CASE__				=							["""TextQuestionAnsweringTool"""]
							SCREAMING_SNAKE_CASE__				=							["""TextSummarizationTool"""]
							SCREAMING_SNAKE_CASE__				=							["""TextToSpeechTool"""]
							SCREAMING_SNAKE_CASE__				=							["""TranslationTool"""]

if TYPE_CHECKING:
							from .agents import Agent, AzureOpenAiAgent, HfAgent, LocalAgent, OpenAiAgent
							from .base import PipelineTool, RemoteTool, Tool, launch_gradio_demo, load_tool

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .document_question_answering import DocumentQuestionAnsweringTool
														from .image_captioning import ImageCaptioningTool
														from .image_question_answering import ImageQuestionAnsweringTool
														from .image_segmentation import ImageSegmentationTool
														from .speech_to_text import SpeechToTextTool
														from .text_classification import TextClassificationTool
														from .text_question_answering import TextQuestionAnsweringTool
														from .text_summarization import TextSummarizationTool
														from .text_to_speech import TextToSpeechTool
														from .translation import TranslationTool
else:
							import sys

							SCREAMING_SNAKE_CASE__				=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





