



"""simple docstring"""






import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_torch, require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_torch_available, is_vision_available


if is_torch_available():
	import torch


if is_vision_available():
	from PIL import Image

	from transformers import MgpstrProcessor, ViTImageProcessor



@require_torch
@require_vision
class     A__     (						unittest.TestCase							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE					   =			ViTImageProcessor if is_vision_available() else None

	@property
	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> int:

				"""simple docstring"""

				return self.image_processor_tester.prepare_image_processor_dict()

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> List[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						List[str]				= (3, 32, 128)
				__lowerCAmelCase   :						Tuple				= tempfile.mkdtemp()

				# fmt: off
				__lowerCAmelCase   :						Tuple				= ["[GO]", "[s]", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"]
				# fmt: on
				__lowerCAmelCase   :						Optional[Any]				= dict(zip(_SCREAMING_SNAKE_CASE ,    range(len(_SCREAMING_SNAKE_CASE))))

				__lowerCAmelCase   :						int				= os.path.join(self.tmpdirname ,    VOCAB_FILES_NAMES["vocab_file"])
				with open(self.vocab_file ,    "w" ,    encoding="utf-8") as fp:
							fp.write(json.dumps(_SCREAMING_SNAKE_CASE) + "\n")

				__lowerCAmelCase   :						Optional[Any]				= {
				    "do_normalize": False,
				    "do_resize": True,
				    "image_processor_type": "ViTImageProcessor",
				    "resample": 3,
				    "size": {"height": 32, "width": 128},
				}
				__lowerCAmelCase   :						Any				= os.path.join(self.tmpdirname ,    _SCREAMING_SNAKE_CASE)
				with open(self.image_processor_file ,    "w" ,    encoding="utf-8") as fp:
							json.dump(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any] ,    **_SCREAMING_SNAKE_CASE:	Union[str, Any])       -> Optional[int]:

				"""simple docstring"""

				return MgpstrTokenizer.from_pretrained(self.tmpdirname ,    **_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    **_SCREAMING_SNAKE_CASE:	int)       -> str:

				"""simple docstring"""

				return ViTImageProcessor.from_pretrained(self.tmpdirname ,    **_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	int)       -> Tuple:

				"""simple docstring"""

				shutil.rmtree(self.tmpdirname)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= np.random.randint(255 ,    size=(3, 30, 400) ,    dtype=np.uinta)

				__lowerCAmelCase   :						Any				= Image.fromarray(np.moveaxis(_SCREAMING_SNAKE_CASE ,    0 ,    -1))

				return image_input

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> List[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						List[str]				= self.get_tokenizer()
				__lowerCAmelCase   :						Union[str, Any]				= self.get_image_processor()

				__lowerCAmelCase   :						Optional[Any]				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)
				processor.save_pretrained(self.tmpdirname)
				__lowerCAmelCase   :						Dict				= MgpstrProcessor.from_pretrained(self.tmpdirname ,    use_fast=_SCREAMING_SNAKE_CASE)

				self.assertEqual(processor.char_tokenizer.get_vocab() ,    tokenizer.get_vocab())
				self.assertIsInstance(processor.char_tokenizer ,    _SCREAMING_SNAKE_CASE)

				self.assertEqual(processor.image_processor.to_json_string() ,    image_processor.to_json_string())
				self.assertIsInstance(processor.image_processor ,    _SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Tuple)       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						Dict				= self.get_tokenizer()
				__lowerCAmelCase   :						List[str]				= self.get_image_processor()

				__lowerCAmelCase   :						Optional[Any]				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)
				processor.save_pretrained(self.tmpdirname)

				__lowerCAmelCase   :						Tuple				= self.get_tokenizer(bos_token="(BOS)" ,    eos_token="(EOS)")
				__lowerCAmelCase   :						Any				= self.get_image_processor(do_normalize=_SCREAMING_SNAKE_CASE ,    padding_value=1.0)

				__lowerCAmelCase   :						Optional[Any]				= MgpstrProcessor.from_pretrained(
				    self.tmpdirname ,    bos_token="(BOS)" ,    eos_token="(EOS)" ,    do_normalize=_SCREAMING_SNAKE_CASE ,    padding_value=1.0)

				self.assertEqual(processor.char_tokenizer.get_vocab() ,    tokenizer_add_kwargs.get_vocab())
				self.assertIsInstance(processor.char_tokenizer ,    _SCREAMING_SNAKE_CASE)

				self.assertEqual(processor.image_processor.to_json_string() ,    image_processor_add_kwargs.to_json_string())
				self.assertIsInstance(processor.image_processor ,    _SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[str])       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[Any]				= self.get_image_processor()
				__lowerCAmelCase   :						Dict				= self.get_tokenizer()

				__lowerCAmelCase   :						Any				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						str				= self.prepare_image_inputs()

				__lowerCAmelCase   :						str				= image_processor(_SCREAMING_SNAKE_CASE ,    return_tensors="np")
				__lowerCAmelCase   :						Union[str, Any]				= processor(images=_SCREAMING_SNAKE_CASE ,    return_tensors="np")

				for key in input_image_proc.keys():
							self.assertAlmostEqual(input_image_proc[key].sum() ,    input_processor[key].sum() ,    delta=1e-2)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						int				= self.get_image_processor()
				__lowerCAmelCase   :						Dict				= self.get_tokenizer()

				__lowerCAmelCase   :						Tuple				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Tuple				= "test"

				__lowerCAmelCase   :						List[Any]				= processor(text=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						List[Any]				= tokenizer(_SCREAMING_SNAKE_CASE)
				for key in encoded_tok.keys():
							self.assertListEqual(encoded_tok[key] ,    encoded_processor[key])

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> List[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						str				= self.get_image_processor()
				__lowerCAmelCase   :						Union[str, Any]				= self.get_tokenizer()

				__lowerCAmelCase   :						Optional[int]				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Any				= "test"
				__lowerCAmelCase   :						List[str]				= self.prepare_image_inputs()

				__lowerCAmelCase   :						Optional[int]				= processor(text=_SCREAMING_SNAKE_CASE ,    images=_SCREAMING_SNAKE_CASE)

				self.assertListEqual(list(inputs.keys()) ,    ["pixel_values", "labels"])

				# test if it raises when no input is passed
				with pytest.raises(_SCREAMING_SNAKE_CASE):
							processor()

	def 							_SCREAMING_SNAKE_CASE     (  self:	Union[str, Any])       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[Any]				= self.get_image_processor()
				__lowerCAmelCase   :						Union[str, Any]				= self.get_tokenizer()

				__lowerCAmelCase   :						Any				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						str				= [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9], [3, 4, 3, 1, 1, 8, 9]]

				__lowerCAmelCase   :						Optional[Any]				= processor.char_decode(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= tokenizer.batch_decode(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						int				= [seq.replace(" " ,    "") for seq in decoded_tok]

				self.assertListEqual(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Tuple)       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= self.get_image_processor()
				__lowerCAmelCase   :						str				= self.get_tokenizer()

				__lowerCAmelCase   :						Union[str, Any]				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Union[str, Any]				= None
				__lowerCAmelCase   :						Optional[int]				= self.prepare_image_inputs()

				__lowerCAmelCase   :						List[Any]				= processor(text=_SCREAMING_SNAKE_CASE ,    images=_SCREAMING_SNAKE_CASE)

				self.assertListEqual(list(inputs.keys()) ,    processor.model_input_names)







	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						int				= self.get_image_processor()
				__lowerCAmelCase   :						Optional[int]				= self.get_tokenizer()

				__lowerCAmelCase   :						List[Any]				= MgpstrProcessor(tokenizer=_SCREAMING_SNAKE_CASE ,    image_processor=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						List[str]				= torch.randn(1 ,    27 ,    38)
				__lowerCAmelCase   :						Optional[int]				= torch.randn(1 ,    27 ,    5_0257)
				__lowerCAmelCase   :						int				= torch.randn(1 ,    27 ,    3_0522)

				__lowerCAmelCase   :						int				= processor.batch_decode([char_input, bpe_input, wp_input])

				self.assertListEqual(list(results.keys()) ,    ["generated_text", "scores", "char_preds", "bpe_preds", "wp_preds"])




"""simple docstring"""






import torch
from torch import nn



class     A__     (						nn.Module							):




	'''simple docstring'''





	def __init__(  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=1 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=False)       -> Dict:

				"""simple docstring"""

				super().__init__()

				__lowerCAmelCase   :						Optional[int]				= n_token
				__lowerCAmelCase   :						int				= d_embed
				__lowerCAmelCase   :						str				= d_proj

				__lowerCAmelCase   :						Optional[Any]				= cutoffs + [n_token]
				__lowerCAmelCase   :						Any				= [0] + self.cutoffs
				__lowerCAmelCase   :						Tuple				= div_val

				__lowerCAmelCase   :						Optional[Any]				= self.cutoffs[0]
				__lowerCAmelCase   :						Union[str, Any]				= len(self.cutoffs) - 1
				__lowerCAmelCase   :						List[str]				= self.shortlist_size + self.n_clusters

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

				__lowerCAmelCase   :						List[Any]				= nn.ModuleList()
				__lowerCAmelCase   :						Any				= 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(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)))
										else:
													self.out_projs.append(_SCREAMING_SNAKE_CASE)

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

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

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

				__lowerCAmelCase   :						str				= keep_order

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	int ,    _SCREAMING_SNAKE_CASE:	Dict)       -> int:

				"""simple docstring"""

				if proj is None:
							__lowerCAmelCase   :						Dict				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
				else:
							# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    proj.t().contiguous())
							__lowerCAmelCase   :						List[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
							# else:
							#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
							#     if bias is not None:
							#         logit = logit + bias

				return logit

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Tuple=None ,    _SCREAMING_SNAKE_CASE:	int=False)       -> List[str]:

				"""simple docstring"""

				if labels is not None:
							# Shift so that tokens < n predict n
							__lowerCAmelCase   :						List[Any]				= hidden[..., :-1, :].contiguous()
							__lowerCAmelCase   :						Tuple				= labels[..., 1:].contiguous()
							__lowerCAmelCase   :						Union[str, Any]				= hidden.view(-1 ,    hidden.size(-1))
							__lowerCAmelCase   :						Optional[int]				= 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:
							__lowerCAmelCase   :						List[Any]				= hidden.view(-1 ,    hidden.size(-1))

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

										if i == 0:
													__lowerCAmelCase   :						int				= torch.cat([weight_i, self.cluster_weight] ,    dim=0)
													__lowerCAmelCase   :						Union[str, Any]				= torch.cat([bias_i, self.cluster_bias] ,    dim=0)

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[0], biases[0], self.out_projs[0]

							__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							if labels is None:
										__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							else:
										__lowerCAmelCase   :						int				= torch.zeros_like(_SCREAMING_SNAKE_CASE ,    dtype=hidden.dtype ,    device=hidden.device)

							__lowerCAmelCase   :						str				= 0
							__lowerCAmelCase   :						List[str]				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if labels is not None:
													__lowerCAmelCase   :						List[str]				= (labels >= l_idx) & (labels < r_idx)
													__lowerCAmelCase   :						Any				= mask_i.nonzero().squeeze()

													if indices_i.numel() == 0:
																continue

													__lowerCAmelCase   :						Optional[Any]				= labels.index_select(0 ,    _SCREAMING_SNAKE_CASE) - l_idx
													__lowerCAmelCase   :						List[Any]				= head_logprob.index_select(0 ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						Tuple				= hidden.index_select(0 ,    _SCREAMING_SNAKE_CASE)
										else:
													__lowerCAmelCase   :						List[str]				= hidden

										if i == 0:
													if labels is not None:
																__lowerCAmelCase   :						Optional[int]				= head_logprob_i.gather(1 ,    target_i[:, None]).squeeze(1)
													else:
																__lowerCAmelCase   :						Optional[Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[str]				= weights[i], biases[i], self.out_projs[i]

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

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

				return out







	def 							_SCREAMING_SNAKE_CASE     (  self:	str ,    _SCREAMING_SNAKE_CASE:	List[Any])       -> Optional[Any]:

				"""simple docstring"""

				if self.n_clusters == 0:
							__lowerCAmelCase   :						int				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    self.out_layers[0].weight ,    self.out_layers[0].bias ,    self.out_projs[0])
							return nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)
				else:
							# construct weights and biases
							__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= [], []
							for i in range(len(self.cutoffs)):
										if self.div_val == 1:
													__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
													__lowerCAmelCase   :						Optional[Any]				= self.out_layers[0].weight[l_idx:r_idx]
													__lowerCAmelCase   :						Union[str, Any]				= self.out_layers[0].bias[l_idx:r_idx]
										else:
													__lowerCAmelCase   :						Dict				= self.out_layers[i].weight
													__lowerCAmelCase   :						List[str]				= self.out_layers[i].bias

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

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= weights[0], biases[0], self.out_projs[0]
							__lowerCAmelCase   :						str				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

							__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							__lowerCAmelCase   :						Union[str, Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							__lowerCAmelCase   :						Dict				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if i == 0:
													__lowerCAmelCase   :						Union[str, Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[i], biases[i], self.out_projs[i]

													__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						List[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

													__lowerCAmelCase   :						int				= head_logprob[:, -i] + tail_logprob_i
													__lowerCAmelCase   :						Tuple				= logprob_i

							return out