

import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin
@dataclass
class 					A_   (      _a						):


    '''simple docstring'''






    a__     							=       42
    a__     							=       42






def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE=0.999       ,   SCREAMING_SNAKE_CASE="cosine"       ,   )      ->     Any:






    '''simple docstring'''
    if alpha_transform_type == "cosine":

        def alpha_bar_fn(SCREAMING_SNAKE_CASE      ):
            return math.cos((t + 0.008) / 1.008 * math.pi / 2      ) ** 2

    elif alpha_transform_type == "exp":

        def alpha_bar_fn(SCREAMING_SNAKE_CASE      ):
            return math.exp(t * -12.0      )

    else:
        raise ValueError(f'''Unsupported alpha_tranform_type: {alpha_transform_type}'''      )

    __UpperCAmelCase									=			[]
    for i in range(SCREAMING_SNAKE_CASE      ):
        __UpperCAmelCase									=			i / num_diffusion_timesteps
        __UpperCAmelCase									=			(i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar_fn(SCREAMING_SNAKE_CASE      ) / alpha_bar_fn(SCREAMING_SNAKE_CASE      )       ,   SCREAMING_SNAKE_CASE      )      )
    return torch.tensor(SCREAMING_SNAKE_CASE       ,   dtype=torch.floataa      )
class 					A_   (      _a ,     _a						):


    '''simple docstring'''






    a__     							=       1


    @register_to_config
    def __init__(self     ,      lowercase__ = 1_000     ,      lowercase__ = 0.0001     ,      lowercase__ = 0.02     ,      lowercase__ = "linear"     ,      lowercase__ = 0.0     ,      lowercase__ = None     ,      lowercase__ = True     ,      )     ->      Tuple:
        if trained_betas is not None:
            __UpperCAmelCase									=			torch.from_numpy(lowercase__  )
        elif beta_schedule == "linear":
            __UpperCAmelCase									=			torch.linspace(lowercase__     ,      lowercase__     ,      lowercase__     ,      dtype=torch.floataa  )
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            __UpperCAmelCase									=			(
                torch.linspace(beta_start**0.5     ,      beta_end**0.5     ,      lowercase__     ,      dtype=torch.floataa  ) ** 2
            )
        elif beta_schedule == "squaredcos_cap_v2":
            # Glide cosine schedule
            __UpperCAmelCase									=			betas_for_alpha_bar(lowercase__  )
        elif beta_schedule == "sigmoid":
            # GeoDiff sigmoid schedule
            __UpperCAmelCase									=			torch.linspace(-6     ,      6     ,      lowercase__  )
            __UpperCAmelCase									=			torch.sigmoid(lowercase__  ) * (beta_end - beta_start) + beta_start
        else:
            raise NotImplementedError(F'''{beta_schedule} does is not implemented for {self.__class__}'''  )

        __UpperCAmelCase									=			1.0 - self.betas
        __UpperCAmelCase									=			torch.cumprod(self.alphas     ,      dim=0  )
        __UpperCAmelCase									=			torch.tensor(1.0  )

        __UpperCAmelCase									=			torch.tensor(1.0  )

        # standard deviation of the initial noise distribution
        __UpperCAmelCase									=			1.0

        # setable values
        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			torch.from_numpy(np.arange(0     ,      lowercase__  )[::-1].copy()  )

        __UpperCAmelCase									=			eta


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__ = None  )     ->      torch.FloatTensor:
        return sample


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__ = 10     ,      lowercase__ = 10     ,      lowercase__ = None     ,      )     ->      Union[str, Any]:
        __UpperCAmelCase									=			min(self.config.num_train_timesteps     ,      lowercase__  )
        __UpperCAmelCase									=			num_inference_steps

        __UpperCAmelCase									=			[]

        __UpperCAmelCase									=			{}
        for j in range(0     ,      num_inference_steps - jump_length     ,      lowercase__  ):
            __UpperCAmelCase									=			jump_n_sample - 1

        __UpperCAmelCase									=			num_inference_steps
        while t >= 1:
            __UpperCAmelCase									=			t - 1
            timesteps.append(lowercase__  )

            if jumps.get(lowercase__     ,      0  ) > 0:
                __UpperCAmelCase									=			jumps[t] - 1
                for _ in range(lowercase__  ):
                    __UpperCAmelCase									=			t + 1
                    timesteps.append(lowercase__  )

        __UpperCAmelCase									=			np.array(lowercase__  ) * (self.config.num_train_timesteps // self.num_inference_steps)
        __UpperCAmelCase									=			torch.from_numpy(lowercase__  ).to(lowercase__  )


    def      lowerCAmelCase_						(self     ,      lowercase__  )     ->      Tuple:
        __UpperCAmelCase									=			t - self.config.num_train_timesteps // self.num_inference_steps

        __UpperCAmelCase									=			self.alphas_cumprod[t]
        __UpperCAmelCase									=			self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
        __UpperCAmelCase									=			1 - alpha_prod_t
        __UpperCAmelCase									=			1 - alpha_prod_t_prev

        # For t > 0, compute predicted variance βt (see formula (6) and (7) from
        # https://arxiv.org/pdf/2006.11239.pdf) and sample from it to get
        # previous sample x_{t-1} ~ N(pred_prev_sample, variance) == add
        # variance to pred_sample
        # Is equivalent to formula (16) in https://arxiv.org/pdf/2010.02502.pdf
        # without eta.
        # variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t]
        __UpperCAmelCase									=			(beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

        return variance


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__ = None     ,      lowercase__ = True     ,      )     ->      Union[RePaintSchedulerOutput, Tuple]:
        __UpperCAmelCase									=			timestep
        __UpperCAmelCase									=			timestep - self.config.num_train_timesteps // self.num_inference_steps

        # 1. compute alphas, betas
        __UpperCAmelCase									=			self.alphas_cumprod[t]
        __UpperCAmelCase									=			self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
        __UpperCAmelCase									=			1 - alpha_prod_t

        # 2. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
        __UpperCAmelCase									=			(sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5

        # 3. Clip "predicted x_0"
        if self.config.clip_sample:
            __UpperCAmelCase									=			torch.clamp(lowercase__     ,      -1     ,      1  )

        # We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
        # substitute formula (7) in the algorithm coming from DDPM paper
        # (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
        # DDIM schedule gives the same results as DDPM with eta = 1.0
        # Noise is being reused in 7. and 8., but no impact on quality has
        # been observed.

        # 5. Add noise
        __UpperCAmelCase									=			model_output.device
        __UpperCAmelCase									=			randn_tensor(model_output.shape     ,      generator=lowercase__     ,      device=lowercase__     ,      dtype=model_output.dtype  )
        __UpperCAmelCase									=			self.eta * self._get_variance(lowercase__  ) ** 0.5

        __UpperCAmelCase									=			0
        if t > 0 and self.eta > 0:
            __UpperCAmelCase									=			std_dev_t * noise

        # 6. compute "direction pointing to x_t" of formula (12)
        # from https://arxiv.org/pdf/2010.02502.pdf
        __UpperCAmelCase									=			(1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * model_output

        # 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        __UpperCAmelCase									=			alpha_prod_t_prev**0.5 * pred_original_sample + pred_sample_direction + variance

        # 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
        __UpperCAmelCase									=			(alpha_prod_t_prev**0.5) * original_image + ((1 - alpha_prod_t_prev) ** 0.5) * noise

        # 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
        __UpperCAmelCase									=			mask * prev_known_part + (1.0 - mask) * prev_unknown_part

        if not return_dict:
            return (
                pred_prev_sample,
                pred_original_sample,
            )

        return RePaintSchedulerOutput(prev_sample=lowercase__     ,      pred_original_sample=lowercase__  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__=None  )     ->      List[Any]:
        __UpperCAmelCase									=			self.config.num_train_timesteps // self.num_inference_steps

        for i in range(lowercase__  ):
            __UpperCAmelCase									=			self.betas[timestep + i]
            if sample.device.type == "mps":
                # randn does not work reproducibly on mps
                __UpperCAmelCase									=			randn_tensor(sample.shape     ,      dtype=sample.dtype     ,      generator=lowercase__  )
                __UpperCAmelCase									=			noise.to(sample.device  )
            else:
                __UpperCAmelCase									=			randn_tensor(sample.shape     ,      generator=lowercase__     ,      device=sample.device     ,      dtype=sample.dtype  )

            # 10. Algorithm 1 Line 10 https://arxiv.org/pdf/2201.09865.pdf
            __UpperCAmelCase									=			(1 - beta) ** 0.5 * sample + beta**0.5 * noise

        return sample


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      torch.FloatTensor:
        raise NotImplementedError('''Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.'''  )






    def __len__(self  )     ->      List[Any]:
        return self.config.num_train_timesteps







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFXLMForMultipleChoice,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForSequenceClassification,
           TFXLMForTokenClassification,
           TFXLMModel,
           TFXLMWithLMHeadModel,
           XLMConfig,
       )
class 					A_   :


    '''simple docstring'''






    def __init__(self     ,      lowercase__     ,      )     ->      Tuple:
        __UpperCAmelCase									=			parent
        __UpperCAmelCase									=			13
        __UpperCAmelCase									=			7
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			False
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			99
        __UpperCAmelCase									=			0
        __UpperCAmelCase									=			32
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			4
        __UpperCAmelCase									=			0.1
        __UpperCAmelCase									=			0.1
        __UpperCAmelCase									=			512
        __UpperCAmelCase									=			16
        __UpperCAmelCase									=			2
        __UpperCAmelCase									=			0.02
        __UpperCAmelCase									=			3
        __UpperCAmelCase									=			4
        __UpperCAmelCase									=			'''last'''
        __UpperCAmelCase									=			True
        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			0


    def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
        __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.vocab_size  )
        __UpperCAmelCase									=			random_attention_mask([self.batch_size, self.seq_length]     ,      dtype=tf.floataa  )

        __UpperCAmelCase									=			None
        if self.use_input_lengths:
            __UpperCAmelCase									=			(
                ids_tensor([self.batch_size]     ,      vocab_size=2  ) + self.seq_length - 2
            )  # small variation of seq_length

        __UpperCAmelCase									=			None
        if self.use_token_type_ids:
            __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.n_langs  )

        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			None
        __UpperCAmelCase									=			None
        if self.use_labels:
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      self.type_sequence_label_size  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size, self.seq_length]     ,      self.num_labels  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      2     ,      dtype=tf.floataa  )
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      self.num_choices  )

        __UpperCAmelCase									=			XLMConfig(
            vocab_size=self.vocab_size     ,      n_special=self.n_special     ,      emb_dim=self.hidden_size     ,      n_layers=self.num_hidden_layers     ,      n_heads=self.num_attention_heads     ,      dropout=self.hidden_dropout_prob     ,      attention_dropout=self.attention_probs_dropout_prob     ,      gelu_activation=self.gelu_activation     ,      sinusoidal_embeddings=self.sinusoidal_embeddings     ,      asm=self.asm     ,      causal=self.causal     ,      n_langs=self.n_langs     ,      max_position_embeddings=self.max_position_embeddings     ,      initializer_range=self.initializer_range     ,      summary_type=self.summary_type     ,      use_proj=self.use_proj     ,      bos_token_id=self.bos_token_id     ,      )

        return (
            config,
            input_ids,
            token_type_ids,
            input_lengths,
            sequence_labels,
            token_labels,
            is_impossible_labels,
            choice_labels,
            input_mask,
        )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      List[Any]:
        __UpperCAmelCase									=			TFXLMModel(config=lowercase__  )
        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )

        __UpperCAmelCase									=			[input_ids, input_mask]
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.last_hidden_state.shape     ,      (self.batch_size, self.seq_length, self.hidden_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      List[str]:
        __UpperCAmelCase									=			TFXLMWithLMHeadModel(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )

        __UpperCAmelCase									=			outputs

        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.seq_length, self.vocab_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Optional[Any]:
        __UpperCAmelCase									=			TFXLMForQuestionAnsweringSimple(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths}

        __UpperCAmelCase									=			model(lowercase__  )

        self.parent.assertEqual(result.start_logits.shape     ,      (self.batch_size, self.seq_length)  )
        self.parent.assertEqual(result.end_logits.shape     ,      (self.batch_size, self.seq_length)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Dict:
        __UpperCAmelCase									=			TFXLMForSequenceClassification(lowercase__  )

        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''lengths''': input_lengths}

        __UpperCAmelCase									=			model(lowercase__  )

        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.type_sequence_label_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      int:
        __UpperCAmelCase									=			self.num_labels
        __UpperCAmelCase									=			TFXLMForTokenClassification(config=lowercase__  )
        __UpperCAmelCase									=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.seq_length, self.num_labels)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      )     ->      Tuple:
        __UpperCAmelCase									=			self.num_choices
        __UpperCAmelCase									=			TFXLMForMultipleChoice(config=lowercase__  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			tf.tile(tf.expand_dims(lowercase__     ,      1  )     ,      (1, self.num_choices, 1)  )
        __UpperCAmelCase									=			{
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        __UpperCAmelCase									=			model(lowercase__  )
        self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_choices)  )






    def      lowerCAmelCase_						(self  )     ->      str:
        __UpperCAmelCase									=			self.prepare_config_and_inputs()
        (
            (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     (
            __UpperCAmelCase
        )					,     
        )									=			config_and_inputs
        __UpperCAmelCase									=			{
            '''input_ids''': input_ids,
            '''token_type_ids''': token_type_ids,
            '''langs''': token_type_ids,
            '''lengths''': input_lengths,
        }
        return config, inputs_dict






@require_tf
class 					A_   (      _a ,     _a ,     unittest.TestCase						):


    '''simple docstring'''






    a__     							=       (
        (
            TFXLMModel,
            TFXLMWithLMHeadModel,
            TFXLMForSequenceClassification,
            TFXLMForQuestionAnsweringSimple,
            TFXLMForTokenClassification,
            TFXLMForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    a__     							=       (
        (TFXLMWithLMHeadModel,) if is_tf_available() else ()
    )  # TODO (PVP): Check other models whether language generation is also applicable
    a__     							=       (
        {
            "feature-extraction": TFXLMModel,
            "fill-mask": TFXLMWithLMHeadModel,
            "question-answering": TFXLMForQuestionAnsweringSimple,
            "text-classification": TFXLMForSequenceClassification,
            "text-generation": TFXLMWithLMHeadModel,
            "token-classification": TFXLMForTokenClassification,
            "zero-shot": TFXLMForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    a__     							=       False
    a__     							=       False


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )     ->      Tuple:
        if (
            pipeline_test_casse_name == "QAPipelineTests"
            and tokenizer_name is not None
            and not tokenizer_name.endswith('''Fast'''  )
        ):
            # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
            # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
            # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
            return True

        return False


    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			TFXLMModelTester(self  )
        __UpperCAmelCase									=			ConfigTester(self     ,      config_class=lowercase__     ,      emb_dim=37  )


    def      lowerCAmelCase_						(self  )     ->      str:
        self.config_tester.run_common_tests()


    def      lowerCAmelCase_						(self  )     ->      Optional[int]:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_model(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Any:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_lm_head(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Dict:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_qa(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      int:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_sequence_classif(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_token_classification(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_multiple_choice(*lowercase__  )






    @slow
    def      lowerCAmelCase_						(self  )     ->      int:
        for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCAmelCase									=			TFXLMModel.from_pretrained(lowercase__  )
            self.assertIsNotNone(lowercase__  )






@require_tf
class 					A_   (      unittest.TestCase						):


    '''simple docstring'''






    @slow
    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase									=			TFXLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''  )
        __UpperCAmelCase									=			tf.convert_to_tensor([[14, 447]]     ,      dtype=tf.intaa  )  # the president
        __UpperCAmelCase									=			[
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
        ]  # the president the president the president the president the president the president the president the president the president the president
        # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
        __UpperCAmelCase									=			model.generate(lowercase__     ,      do_sample=lowercase__  )
        self.assertListEqual(output_ids[0].numpy().tolist()     ,      lowercase__  )




