




from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
		from transformers import (
		    TFBertModel,
		    TFCLIPVisionModel,
		    TFDeiTModel,
		    TFRobertaModel,
		    TFVisionTextDualEncoderModel,
		    TFViTModel,
		    VisionTextDualEncoderConfig,
		)

if is_vision_available():
		from PIL import Image

		from transformers import VisionTextDualEncoderProcessor




def 			__magic_name__				( A		:					str   ):






			'''simple docstring'''







			if isinstance(A,  collections.abc.Iterable   ):
						return x
			return (x, x)




@require_tf
class     snake_case__							:
				"""simple docstring"""



				def 							__UpperCAmelCase							(      self				:     int					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     List[Any]							)					->     Optional[Any]:
							pass





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





				def 							__UpperCAmelCase							(      self				:     int							)					->     str:
							pass





				def 							__UpperCAmelCase							(      self				:     Optional[Any]					,      __lowerCamelCase				:     Dict					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     List[Any]=None					,      **__lowerCamelCase				:     List[Any]							)					->     List[str]:
							a    =			VisionTextDualEncoderConfig.from_vision_text_configs(__lowerCamelCase					,      __lowerCamelCase							)

							a    =			TFVisionTextDualEncoderModel(__lowerCamelCase							)

							a    =			model(input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase							)

							self.assertEqual(output["text_embeds"].shape					,      (input_ids.shape[0], config.projection_dim)							)
							self.assertEqual(output["image_embeds"].shape					,      (pixel_values.shape[0], config.projection_dim)							)





				def 							__UpperCAmelCase							(      self				:     Optional[int]					,      __lowerCamelCase				:     int					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     Optional[Any]					,      __lowerCamelCase				:     List[Any]=None					,      **__lowerCamelCase				:     str							)					->     List[Any]:
							a    ,		a    =			self.get_vision_text_model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase					,      text_model=__lowerCamelCase							)

							a    =			model(input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase							)

							self.assertEqual(output["text_embeds"].shape					,      (input_ids.shape[0], model.config.projection_dim)							)
							self.assertEqual(output["image_embeds"].shape					,      (pixel_values.shape[0], model.config.projection_dim)							)





				def 							__UpperCAmelCase							(      self				:     int					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     int=None					,      **__lowerCamelCase				:     Tuple							)					->     Union[str, Any]:
							a    ,		a    =			self.get_vision_text_model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			{"vision_model": vision_model, "text_model": text_model}
							a    =			TFVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowerCamelCase							)

							a    =			model(input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase							)

							self.assertEqual(output["text_embeds"].shape					,      (input_ids.shape[0], model.config.projection_dim)							)
							self.assertEqual(output["image_embeds"].shape					,      (pixel_values.shape[0], model.config.projection_dim)							)





				def 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Optional[int]=None					,      **__lowerCamelCase				:     Tuple							)					->     List[Any]:
							a    ,		a    =			self.get_vision_text_model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase					,      text_model=__lowerCamelCase							)

							a    =			model(input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase							)
							a    =			output[0].numpy()

							with tempfile.TemporaryDirectory() as tmpdirname:
										model.save_pretrained(__lowerCamelCase							)
										a    =			TFVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase							)

										a    =			model(input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase							)
										a    =			after_output[0].numpy()
										a    =			np.amax(np.abs(out_a - out_a							)							)
										self.assertLessEqual(__lowerCamelCase					,      1e-5							)





				def 							__UpperCAmelCase							(      self				:     Dict					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     Optional[int]					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     List[Any]=None					,      **__lowerCamelCase				:     Optional[int]							)					->     Union[str, Any]:
							a    ,		a    =			self.get_vision_text_model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase					,      text_model=__lowerCamelCase							)

							a    =			model(
							    input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase					,      output_attentions=__lowerCamelCase							)

							a    =			output.vision_model_output.attentions
							self.assertEqual(len(__lowerCamelCase							)					,      vision_config.num_hidden_layers							)

							# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
							a    =			to_atuple(vision_model.config.image_size							)
							a    =			to_atuple(vision_model.config.patch_size							)
							a    =			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
							a    =			num_patches + 1
							self.assertEqual(vision_attentions[0].shape[-3:]					,      (vision_config.num_attention_heads, seq_len, seq_len)							)

							a    =			output.text_model_output.attentions
							self.assertEqual(len(__lowerCamelCase							)					,      text_config.num_hidden_layers							)

							self.assertEqual(
							    text_attentions[0].shape[-3:]					,      (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])					,      )





				def 							__UpperCAmelCase							(      self				:     Tuple					,      __lowerCamelCase				:     np.ndarray					,      __lowerCamelCase				:     np.ndarray					,      __lowerCamelCase				:     float							)					->     str:
							a    =			np.abs((a - b)							).max()
							self.assertLessEqual(__lowerCamelCase					,      __lowerCamelCase					,      f"""Difference between torch and flax is {diff} (>= {tol})."""							)





				def 							__UpperCAmelCase							(      self				:     int							)					->     Any:
							a    =			self.prepare_config_and_inputs()
							self.check_vision_text_dual_encoder_model(**__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     int							)					->     Dict:
							a    =			self.prepare_config_and_inputs()
							self.check_model_from_pretrained_configs(**__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Any							)					->     Optional[Any]:
							a    =			self.prepare_config_and_inputs()
							self.check_vision_text_dual_encoder_from_pretrained(**__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Any							)					->     Union[str, Any]:
							a    =			self.prepare_config_and_inputs()
							self.check_save_load(**__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     Tuple:
							a    =			self.prepare_config_and_inputs()
							self.check_vision_text_output_attention(**__lowerCamelCase							)





				@slow
				def 							__UpperCAmelCase							(      self				:     List[Any]							)					->     Tuple:
							a    ,		a    =			self.get_pretrained_model_and_inputs()

							a    =			model_a(**__lowerCamelCase							)
							a    =			outputs[0].numpy()

							with tempfile.TemporaryDirectory() as tmp_dirname:
										model_a.save_pretrained(__lowerCamelCase							)
										a    =			TFVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase							)

										a    =			model_a(**__lowerCamelCase							)
										a    =			after_outputs[0].numpy()
										a    =			np.amax(np.abs(out_a - out_a							)							)
										self.assertLessEqual(__lowerCamelCase					,      1e-5							)




@require_tf
class     snake_case__							(_UpperCamelCase    ,       unittest.TestCase   ):
				"""simple docstring"""



				def 							__UpperCAmelCase							(      self				:     str							)					->     Union[str, Any]:
							a    =			TFVisionTextDualEncoderModel.from_vision_text_pretrained(
							    "hf-internal-testing/tiny-random-vit"					,      "hf-internal-testing/tiny-random-bert"							)
							a    =			13
							a    =			floats_tensor(
							    [
							        batch_size,
							        model.vision_model.config.num_channels,
							        model.vision_model.config.image_size,
							        model.vision_model.config.image_size,
							    ]							)
							a    =			ids_tensor([batch_size, 4]					,      model.text_model.config.vocab_size							)
							a    =			random_attention_mask([batch_size, 4]							)
							a    =			{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

							return model, inputs





				def 							__UpperCAmelCase							(      self				:     List[Any]					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     Union[str, Any]							)					->     int:
							a    =			TFViTModel(__lowerCamelCase					,      name="vision_model"							)
							a    =			TFBertModel(__lowerCamelCase					,      name="text_model"							)
							return vision_model, text_model





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     List[Any]:
							a    =			TFViTModelTester(self							)
							a    =			TFBertModelTester(self							)
							a    =			vit_model_tester.prepare_config_and_inputs()
							a    =			bert_model_tester.prepare_config_and_inputs()

							a    ,		a    ,		a    =			vision_config_and_inputs

							(
							    (
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		
							)    =			text_config_and_inputs

							return {
							    "text_config": text_config,
							    "vision_config": vision_config,
							    "pixel_values": pixel_values,
							    "attention_mask": input_mask,
							    "input_ids": input_ids,
							    "text_token_type_ids": token_type_ids,
							    "text_sequence_labels": sequence_labels,
							    "text_token_labels": token_labels,
							    "text_choice_labels": choice_labels,
							}




@require_tf
class     snake_case__							(_UpperCamelCase    ,       unittest.TestCase   ):
				"""simple docstring"""



				def 							__UpperCAmelCase							(      self				:     Any							)					->     List[str]:
							# DeiT repo doesn't have TF weights, but we don't actually use the weights at all so let's
							# just reinitialize it.
							a    =			TFVisionTextDualEncoderModel.from_vision_text_pretrained(
							    "Rocketknight1/tiny-random-deit-tf"					,      "hf-internal-testing/tiny-random-roberta"							)
							a    =			13
							a    =			floats_tensor(
							    [
							        batch_size,
							        model.vision_model.config.num_channels,
							        model.vision_model.config.image_size,
							        model.vision_model.config.image_size,
							    ]							)
							a    =			ids_tensor([batch_size, 4]					,      model.text_model.config.vocab_size							)
							a    =			random_attention_mask([batch_size, 4]							)
							a    =			{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

							return model, inputs





				def 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     int					,      __lowerCamelCase				:     Any					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     List[str]=None					,      **__lowerCamelCase				:     List[Any]							)					->     Dict:
							a    ,		a    =			self.get_vision_text_model(__lowerCamelCase					,      __lowerCamelCase							)
							a    =			TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase					,      text_model=__lowerCamelCase							)

							a    =			model(
							    input_ids=__lowerCamelCase					,      pixel_values=__lowerCamelCase					,      attention_mask=__lowerCamelCase					,      output_attentions=__lowerCamelCase							)

							a    =			output.vision_model_output.attentions
							self.assertEqual(len(__lowerCamelCase							)					,      vision_config.num_hidden_layers							)

							# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
							a    =			to_atuple(vision_model.config.image_size							)
							a    =			to_atuple(vision_model.config.patch_size							)
							a    =			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
							a    =			num_patches + 2
							self.assertEqual(vision_attentions[0].shape[-3:]					,      (vision_config.num_attention_heads, seq_len, seq_len)							)

							a    =			output.text_model_output.attentions
							self.assertEqual(len(__lowerCamelCase							)					,      text_config.num_hidden_layers							)

							self.assertEqual(
							    text_attentions[0].shape[-3:]					,      (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])					,      )





				def 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     Optional[Any]					,      __lowerCamelCase				:     Optional[int]							)					->     Any:
							a    =			TFDeiTModel(__lowerCamelCase					,      name="vision_model"							)
							a    =			TFRobertaModel(__lowerCamelCase					,      name="text_model"							)
							return vision_model, text_model





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]							)					->     Union[str, Any]:
							a    =			TFDeiTModelTester(self							)
							a    =			TFRobertaModelTester(self							)
							a    =			vit_model_tester.prepare_config_and_inputs()
							a    =			bert_model_tester.prepare_config_and_inputs()

							a    ,		a    ,		a    =			vision_config_and_inputs

							(
							    (
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		
							)    =			text_config_and_inputs

							return {
							    "text_config": text_config,
							    "vision_config": vision_config,
							    "pixel_values": pixel_values,
							    "attention_mask": input_mask,
							    "input_ids": input_ids,
							    "text_token_type_ids": token_type_ids,
							    "text_sequence_labels": sequence_labels,
							    "text_token_labels": token_labels,
							    "text_choice_labels": choice_labels,
							}




@require_tf
class     snake_case__							(_UpperCamelCase    ,       unittest.TestCase   ):
				"""simple docstring"""



				def 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     Optional[int]:
							a    =			TFVisionTextDualEncoderModel.from_vision_text_pretrained(
							    "Rocketknight1/tiny-random-clip-tf"					,      "hf-internal-testing/tiny-random-bert"							)
							a    =			13
							a    =			floats_tensor(
							    [
							        batch_size,
							        model.vision_model.config.num_channels,
							        model.vision_model.config.image_size,
							        model.vision_model.config.image_size,
							    ]							)
							a    =			ids_tensor([batch_size, 4]					,      model.text_model.config.vocab_size							)
							a    =			random_attention_mask([batch_size, 4]							)
							a    =			{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

							return model, inputs





				def 							__UpperCAmelCase							(      self				:     List[Any]					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     List[str]							)					->     int:
							a    =			TFCLIPVisionModel(__lowerCamelCase					,      name="vision_model"							)
							a    =			TFBertModel(__lowerCamelCase					,      name="text_model"							)
							return vision_model, text_model





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Dict:
							a    =			TFCLIPVisionModelTester(self							)
							a    =			TFBertModelTester(self							)
							a    =			clip_model_tester.prepare_config_and_inputs()
							a    =			bert_model_tester.prepare_config_and_inputs()

							a    ,		a    =			vision_config_and_inputs

							(
							    (
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		(
							    a
							)    ,		
							)    =			text_config_and_inputs

							return {
							    "text_config": text_config,
							    "vision_config": vision_config,
							    "pixel_values": pixel_values,
							    "attention_mask": input_mask,
							    "input_ids": input_ids,
							    "text_token_type_ids": token_type_ids,
							    "text_sequence_labels": sequence_labels,
							    "text_token_labels": token_labels,
							    "text_choice_labels": choice_labels,
							}




@require_vision
@require_tf
class     snake_case__							(unittest.TestCase   ):
				"""simple docstring"""



				@slow
				def 							__UpperCAmelCase							(      self				:     List[str]							)					->     Union[str, Any]:
							a    =			TFVisionTextDualEncoderModel.from_pretrained(
							    "clip-italian/clip-italian"					,      logit_scale_init_value=1.0					,      from_pt=__lowerCamelCase							)
							a    =			VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"							)

							a    =			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
							a    =			processor(
							    text=["una foto di un gatto", "una foto di un cane"]					,      images=__lowerCamelCase					,      padding=__lowerCamelCase					,      return_tensors="np"							)

							a    =			model(**__lowerCamelCase							)

							# verify the logits
							self.assertEqual(outputs.logits_per_image.shape					,      (inputs.pixel_values.shape[0], inputs.input_ids.shape[0])							)
							self.assertEqual(
							    outputs.logits_per_text.shape					,      (inputs.input_ids.shape[0], inputs.pixel_values.shape[0])					,      )

							a    =			np.array([[1.2_284_727, 0.3_104_122]]							)

							self.assertTrue(np.allclose(outputs.logits_per_image.numpy()					,      __lowerCamelCase					,      atol=1e-3							)							)








"""simple docstring"""





import torch
from torch import nn






class     _UpperCAmelCase(    nn.Module	):



		def __init__(					self				,		__a				,		__a				,		__a				,		__a				,		__a=1				,		__a=False)       ->       Tuple:






							'''simple docstring'''
							super().__init__()

							_UpperCamelCase							 =  n_token
							_UpperCamelCase							 =  d_embed
							_UpperCamelCase							 =  d_proj

							_UpperCamelCase							 =  cutoffs + [n_token]
							_UpperCamelCase							 =  [0] + self.cutoffs
							_UpperCamelCase							 =  div_val

							_UpperCamelCase							 =  self.cutoffs[0]
							_UpperCamelCase							 =  len(self.cutoffs) - 1
							_UpperCamelCase							 =  self.shortlist_size + self.n_clusters

							if self.n_clusters > 0:
												_UpperCamelCase							 =  nn.Parameter(torch.zeros(self.n_clusters				,		self.d_embed))
												_UpperCamelCase							 =  nn.Parameter(torch.zeros(self.n_clusters))

							_UpperCamelCase							 =  nn.ModuleList()
							_UpperCamelCase							 =  nn.ParameterList()

							if div_val == 1:
												for i in range(len(self.cutoffs)):
																	if d_proj != d_embed:
																						self.out_projs.append(nn.Parameter(torch.FloatTensor(__a				,		__a)))
																	else:
																						self.out_projs.append(__a)

												self.out_layers.append(nn.Linear(__a				,		__a))
							else:
												for i in range(len(self.cutoffs)):
																	_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																	_UpperCamelCase							 =  d_embed // (div_val**i)

																	self.out_projs.append(nn.Parameter(torch.FloatTensor(__a				,		__a)))

																	self.out_layers.append(nn.Linear(__a				,		r_idx - l_idx))

							_UpperCamelCase							 =  keep_order



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






							'''simple docstring'''
							if proj is None:
												_UpperCamelCase							 =  nn.functional.linear(__a				,		__a				,		bias=__a)
							else:
												# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
												_UpperCamelCase							 =  nn.functional.linear(__a				,		proj.t().contiguous())
												_UpperCamelCase							 =  nn.functional.linear(__a				,		__a				,		bias=__a)
												# else:
												#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
												#     if bias is not None:
												#         logit = logit + bias

							return logit



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






							'''simple docstring'''

							if labels is not None:
												# Shift so that tokens < n predict n
												_UpperCamelCase							 =  hidden[..., :-1, :].contiguous()
												_UpperCamelCase							 =  labels[..., 1:].contiguous()
												_UpperCamelCase							 =  hidden.view(-1				,		hidden.size(-1))
												_UpperCamelCase							 =  labels.view(-1)
												if hidden.size(0) != labels.size(0):
																	raise RuntimeError('''Input and labels should have the same size in the batch dimension.''')
							else:
												_UpperCamelCase							 =  hidden.view(-1				,		hidden.size(-1))

							if self.n_clusters == 0:
												_UpperCamelCase							 =  self._compute_logit(__a				,		self.out_layers[0].weight				,		self.out_layers[0].bias				,		self.out_projs[0])
												if labels is not None:
																	_UpperCamelCase							 =  labels != -1_00
																	_UpperCamelCase							 =  torch.zeros_like(__a				,		dtype=hidden.dtype				,		device=hidden.device)
																	_UpperCamelCase							 =  (
																	    -nn.functional.log_softmax(__a				,		dim=-1)[mask].gather(1				,		labels[mask].unsqueeze(1)).squeeze(1)
																	)
												else:
																	_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=-1)
							else:
												# construct weights and biases
												_UpperCamelCase ,  _UpperCamelCase							 =  [], []
												for i in range(len(self.cutoffs)):
																	if self.div_val == 1:
																						_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						_UpperCamelCase							 =  self.out_layers[0].weight[l_idx:r_idx]
																						_UpperCamelCase							 =  self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						_UpperCamelCase							 =  self.out_layers[i].weight
																						_UpperCamelCase							 =  self.out_layers[i].bias

																	if i == 0:
																						_UpperCamelCase							 =  torch.cat([weight_i, self.cluster_weight]				,		dim=0)
																						_UpperCamelCase							 =  torch.cat([bias_i, self.cluster_bias]				,		dim=0)

																	weights.append(__a)
																	biases.append(__a)

												_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[0], biases[0], self.out_projs[0]

												_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
												_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

												if labels is None:
																	_UpperCamelCase							 =  hidden.new_empty((head_logit.size(0), self.n_token))
												else:
																	_UpperCamelCase							 =  torch.zeros_like(__a				,		dtype=hidden.dtype				,		device=hidden.device)

												_UpperCamelCase							 =  0
												_UpperCamelCase							 =  [0] + self.cutoffs
												for i in range(len(__a) - 1):
																	_UpperCamelCase ,  _UpperCamelCase							 =  cutoff_values[i], cutoff_values[i + 1]

																	if labels is not None:
																						_UpperCamelCase							 =  (labels >= l_idx) & (labels < r_idx)
																						_UpperCamelCase							 =  mask_i.nonzero().squeeze()

																						if indices_i.numel() == 0:
																											continue

																						_UpperCamelCase							 =  labels.index_select(0				,		__a) - l_idx
																						_UpperCamelCase							 =  head_logprob.index_select(0				,		__a)
																						_UpperCamelCase							 =  hidden.index_select(0				,		__a)
																	else:
																						_UpperCamelCase							 =  hidden

																	if i == 0:
																						if labels is not None:
																											_UpperCamelCase							 =  head_logprob_i.gather(1				,		target_i[:, None]).squeeze(1)
																						else:
																											_UpperCamelCase							 =  head_logprob[:, : self.cutoffs[0]]
																	else:
																						_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[i], biases[i], self.out_projs[i]

																						_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
																						_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)
																						_UpperCamelCase							 =  self.cutoffs[0] + i - 1  # No probability for the head cluster
																						if labels is not None:
																											_UpperCamelCase							 =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																											    1				,		target_i[:, None]).squeeze(1)
																						else:
																											_UpperCamelCase							 =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																											_UpperCamelCase							 =  logprob_i

																	if labels is not None:
																						if (hasattr(self				,		'''keep_order''') and self.keep_order) or keep_order:
																											out.index_copy_(0				,		__a				,		-logprob_i)
																						else:
																											out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
																						offset += logprob_i.size(0)

							return out




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






							'''simple docstring'''
							if self.n_clusters == 0:
												_UpperCamelCase							 =  self._compute_logit(__a				,		self.out_layers[0].weight				,		self.out_layers[0].bias				,		self.out_projs[0])
												return nn.functional.log_softmax(__a				,		dim=-1)
							else:
												# construct weights and biases
												_UpperCamelCase ,  _UpperCamelCase							 =  [], []
												for i in range(len(self.cutoffs)):
																	if self.div_val == 1:
																						_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						_UpperCamelCase							 =  self.out_layers[0].weight[l_idx:r_idx]
																						_UpperCamelCase							 =  self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						_UpperCamelCase							 =  self.out_layers[i].weight
																						_UpperCamelCase							 =  self.out_layers[i].bias

																	if i == 0:
																						_UpperCamelCase							 =  torch.cat([weight_i, self.cluster_weight]				,		dim=0)
																						_UpperCamelCase							 =  torch.cat([bias_i, self.cluster_bias]				,		dim=0)

																	weights.append(__a)
																	biases.append(__a)

												_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[0], biases[0], self.out_projs[0]
												_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)

												_UpperCamelCase							 =  hidden.new_empty((head_logit.size(0), self.n_token))
												_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

												_UpperCamelCase							 =  [0] + self.cutoffs
												for i in range(len(__a) - 1):
																	_UpperCamelCase ,  _UpperCamelCase							 =  cutoff_values[i], cutoff_values[i + 1]

																	if i == 0:
																						_UpperCamelCase							 =  head_logprob[:, : self.cutoffs[0]]
																	else:
																						_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[i], biases[i], self.out_projs[i]

																						_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
																						_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

																						_UpperCamelCase							 =  head_logprob[:, -i] + tail_logprob_i
																						_UpperCamelCase							 =  logprob_i

												return out






