

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


  """simple docstring"""



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

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

      SCREAMING_SNAKE_CASE__					     =						None
      SCREAMING_SNAKE_CASE__					     =						decoder_seq_length
      SCREAMING_SNAKE_CASE__					     =						2
      SCREAMING_SNAKE_CASE__					     =						1



  def      _UpperCamelCase    (	self    :				Tuple	)					-> Dict:
      SCREAMING_SNAKE_CASE__					     =						ids_tensor([self.batch_size, self.decoder_seq_length]     ,  self.vocab_size	)

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

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

      SCREAMING_SNAKE_CASE__					     =						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      _UpperCamelCase    (	self    :				List[Any]     ,  A_    :				Any     ,  A_    :				int     ,  A_    :				List[Any]     ,  A_    :				List[str]     ,  )					-> Optional[Any]:
      SCREAMING_SNAKE_CASE__					     =						True
      SCREAMING_SNAKE_CASE__					     =						SpeechaTextaDecoder(config=A_	).to(A_	).eval()
      SCREAMING_SNAKE_CASE__					     =						input_ids[:2]

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

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

      SCREAMING_SNAKE_CASE__					     =						outputs["past_key_values"]

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

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

      print(A_	)

      SCREAMING_SNAKE_CASE__					     =						model(A_	)["last_hidden_state"]
      SCREAMING_SNAKE_CASE__					     =						model(A_     ,  past_key_values=A_	)["last_hidden_state"]

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

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







  def      _UpperCamelCase    (	self    :				List[str]	)					-> Union[str, Any]:
      SCREAMING_SNAKE_CASE__					     =						self.prepare_config_and_inputs()
      (
          (
          SCREAMING_SNAKE_CASE__
      )     ,		(
          SCREAMING_SNAKE_CASE__
      )     ,		(
          SCREAMING_SNAKE_CASE__
      )     ,		(
          SCREAMING_SNAKE_CASE__
      )     ,		
      )					     =						config_and_inputs

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







@require_torch
class    UpperCamelCase_			(  _snake_case		,     _snake_case		,     _snake_case		,     unittest.TestCase			):


  """simple docstring"""
  SCREAMING_SNAKE_CASE__				: Optional[int]										=      (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE__				: str										=      (SpeechaTextaForCausalLM,) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE__				: Dict										=      {"""text-generation""": SpeechaTextaForCausalLM} if is_torch_available() else {}
  SCREAMING_SNAKE_CASE__				: Optional[int]										=      True
  SCREAMING_SNAKE_CASE__				: Any										=      False



  def      _UpperCamelCase    (	self    :				Tuple     ,  )					-> Optional[Any]:
      SCREAMING_SNAKE_CASE__					     =						SpeechaTextaStandaloneDecoderModelTester(self     ,  is_training=A_	)
      SCREAMING_SNAKE_CASE__					     =						ConfigTester(self     ,  config_class=A_	)



  def      _UpperCamelCase    (	self    :				List[str]	)					-> List[str]:
      pass



  def      _UpperCamelCase    (	self    :				List[str]	)					-> Dict:
      pass



  def      _UpperCamelCase    (	self    :				int	)					-> Optional[int]:
      pass



  def      _UpperCamelCase    (	self    :				str	)					-> Dict:
      self.config_tester.run_common_tests()



  def      _UpperCamelCase    (	self    :				Optional[Any]	)					-> Any:
      SCREAMING_SNAKE_CASE__					     =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past(*A_	)







  def      _UpperCamelCase    (	self    :				List[Any]	)					-> str:
      return





import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_splinter import SplinterTokenizer


A__        =			logging.get_logger(__name__)

A__        =			{"vocab_file": "vocab.txt"}

A__        =			{
    "vocab_file": {
        "tau/splinter-base": "https://huggingface.co/tau/splinter-base/resolve/main/vocab.txt",
        "tau/splinter-base-qass": "https://huggingface.co/tau/splinter-base-qass/resolve/main/vocab.txt",
        "tau/splinter-large": "https://huggingface.co/tau/splinter-large/resolve/main/vocab.txt",
        "tau/splinter-large-qass": "https://huggingface.co/tau/splinter-large-qass/resolve/main/vocab.txt",
    }
}

A__        =			{
    "tau/splinter-base": 512,
    "tau/splinter-base-qass": 512,
    "tau/splinter-large": 512,
    "tau/splinter-large-qass": 512,
}

A__        =			{
    "tau/splinter-base": {"do_lower_case": False},
    "tau/splinter-base-qass": {"do_lower_case": False},
    "tau/splinter-large": {"do_lower_case": False},
    "tau/splinter-large-qass": {"do_lower_case": False},
}







class    UpperCamelCase_			(  _snake_case			):


  """simple docstring"""
  SCREAMING_SNAKE_CASE__				: Dict										=      VOCAB_FILES_NAMES
  SCREAMING_SNAKE_CASE__				: Optional[int]										=      PRETRAINED_VOCAB_FILES_MAP
  SCREAMING_SNAKE_CASE__				: List[Any]										=      PRETRAINED_INIT_CONFIGURATION
  SCREAMING_SNAKE_CASE__				: str										=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  SCREAMING_SNAKE_CASE__				: Optional[Any]										=      SplinterTokenizer



  def __init__(	self    :				Tuple     ,  A_    :				Optional[Any]=None     ,  A_    :				Optional[int]=None     ,  A_    :				Optional[Any]=True     ,  A_    :				str="[UNK]"     ,  A_    :				Optional[Any]="[SEP]"     ,  A_    :				str="[PAD]"     ,  A_    :				List[Any]="[CLS]"     ,  A_    :				Union[str, Any]="[MASK]"     ,  A_    :				Any="[QUESTION]"     ,  A_    :				List[Any]=True     ,  A_    :				Union[str, Any]=None     ,  **A_    :				Optional[Any]     ,  )					-> Optional[Any]:
      super().__init__(
          A_     ,  tokenizer_file=A_     ,  do_lower_case=A_     ,  unk_token=A_     ,  sep_token=A_     ,  pad_token=A_     ,  cls_token=A_     ,  mask_token=A_     ,  tokenize_chinese_chars=A_     ,  strip_accents=A_     ,  additional_special_tokens=(question_token,)     ,  **A_     ,  )

      SCREAMING_SNAKE_CASE__					     =						json.loads(self.backend_tokenizer.normalizer.__getstate__()	)
      if (
          pre_tok_state.get("lowercase"     ,  A_	) != do_lower_case
          or pre_tok_state.get("strip_accents"     ,  A_	) != strip_accents
      ):
          SCREAMING_SNAKE_CASE__					     =						getattr(A_     ,  pre_tok_state.pop("type"	)	)
          SCREAMING_SNAKE_CASE__					     =						do_lower_case
          SCREAMING_SNAKE_CASE__					     =						strip_accents
          SCREAMING_SNAKE_CASE__					     =						pre_tok_class(**A_	)

      SCREAMING_SNAKE_CASE__					     =						do_lower_case



  @property
  def      _UpperCamelCase    (	self    :				List[Any]	)					-> Optional[Any]:
      return self.convert_tokens_to_ids(self.question_token	)



  def      _UpperCamelCase    (	self    :				List[str]     ,  A_    :				List[int]     ,  A_    :				Optional[List[int]] = None	)					-> List[int]:
      if token_ids_a is None:
          return [self.cls_token_id] + token_ids_a + [self.sep_token_id]

      SCREAMING_SNAKE_CASE__					     =						[self.cls_token_id]
      SCREAMING_SNAKE_CASE__					     =						[self.sep_token_id]
      SCREAMING_SNAKE_CASE__					     =						[self.question_token_id] + [self.convert_tokens_to_ids("."	)]
      if self.padding_side == "right":
          # Input is question-then-context
          return cls + token_ids_a + question_suffix + sep + token_ids_a + sep
      else:
          # Input is context-then-question
          return cls + token_ids_a + sep + token_ids_a + question_suffix + sep



  def      _UpperCamelCase    (	self    :				Optional[int]     ,  A_    :				List[int]     ,  A_    :				Optional[List[int]] = None	)					-> List[int]:
      SCREAMING_SNAKE_CASE__					     =						[self.sep_token_id]
      SCREAMING_SNAKE_CASE__					     =						[self.cls_token_id]
      SCREAMING_SNAKE_CASE__					     =						[self.question_token_id] + [self.convert_tokens_to_ids("."	)]
      if token_ids_a is None:
          return len(cls + token_ids_a + sep	) * [0]

      if self.padding_side == "right":
          # Input is question-then-context
          return len(cls + token_ids_a + question_suffix + sep	) * [0] + len(token_ids_a + sep	) * [1]
      else:
          # Input is context-then-question
          return len(cls + token_ids_a + sep	) * [0] + len(token_ids_a + question_suffix + sep	) * [1]







  def      _UpperCamelCase    (	self    :				Optional[int]     ,  A_    :				str     ,  A_    :				Optional[str] = None	)					-> Tuple[str]:
      SCREAMING_SNAKE_CASE__					     =						self._tokenizer.model.save(A_     ,  name=A_	)
      return tuple(A_	)


