"""simple docstring"""







import torch
from torch import nn







class        UpperCamelCase		(			nn.Module  ):
   def __init__(			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase=1     ,__UpperCamelCase=False			)	->      Dict:


        '''simple docstring'''


        super().__init__()

        lowercase_	:				str													=	n_token
        lowercase_	:				int													=	d_embed
        lowercase_	:				Tuple													=	d_proj

        lowercase_	:				Any													=	cutoffs + [n_token]
        lowercase_	:				Dict													=	[0] + self.cutoffs
        lowercase_	:				Tuple													=	div_val

        lowercase_	:				List[Any]													=	self.cutoffs[0]
        lowercase_	:				Any													=	len(self.cutoffs			) - 1
        lowercase_	:				Optional[Any]													=	self.shortlist_size + self.n_clusters

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

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

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

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

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

        lowercase_	:				List[Any]													=	keep_order
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      int:


        '''simple docstring'''


        if proj is None:
             lowercase_	:				List[Any]													=	nn.functional.linear(__UpperCamelCase     ,__UpperCamelCase     ,bias=__UpperCamelCase			)
        else:
             # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
             lowercase_	:				int													=	nn.functional.linear(__UpperCamelCase     ,proj.t().contiguous()			)
             lowercase_	:				List[Any]													=	nn.functional.linear(__UpperCamelCase     ,__UpperCamelCase     ,bias=__UpperCamelCase			)
             # 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     ,__UpperCamelCase     ,__UpperCamelCase=None     ,__UpperCamelCase=False			)	->      Optional[int]:


        '''simple docstring'''


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

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

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

                  weights.append(__UpperCamelCase			)
                  biases.append(__UpperCamelCase			)

             lowercase_							,  lowercase_							,  lowercase_	:				Union[str, Any]													=	weights[0], biases[0], self.out_projs[0]

             lowercase_	:				Union[str, Any]													=	self._compute_logit(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)
             lowercase_	:				Tuple													=	nn.functional.log_softmax(__UpperCamelCase     ,dim=1			)

             if labels is None:
                  lowercase_	:				List[str]													=	hidden.new_empty((head_logit.size(0			), self.n_token)			)
             else:
                  lowercase_	:				List[Any]													=	torch.zeros_like(__UpperCamelCase     ,dtype=hidden.dtype     ,device=hidden.device			)

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

                  if labels is not None:
                       lowercase_	:				Union[str, Any]													=	(labels >= l_idx) & (labels < r_idx)
                       lowercase_	:				Optional[int]													=	mask_i.nonzero().squeeze()

                       if indices_i.numel() == 0:
                            continue

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

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

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

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

        return out







   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase			)	->      Any:


        '''simple docstring'''


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

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

                  weights.append(__UpperCamelCase			)
                  biases.append(__UpperCamelCase			)

             lowercase_							,  lowercase_							,  lowercase_	:				Any													=	weights[0], biases[0], self.out_projs[0]
             lowercase_	:				Union[str, Any]													=	self._compute_logit(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)

             lowercase_	:				int													=	hidden.new_empty((head_logit.size(0			), self.n_token)			)
             lowercase_	:				List[Any]													=	nn.functional.log_softmax(__UpperCamelCase     ,dim=1			)

             lowercase_	:				int													=	[0] + self.cutoffs
             for i in range(len(__UpperCamelCase			) - 1			):
                  lowercase_							,  lowercase_	:				Optional[int]													=	cutoff_values[i], cutoff_values[i + 1]

                  if i == 0:
                       lowercase_	:				List[str]													=	head_logprob[:, : self.cutoffs[0]]
                  else:
                       lowercase_							,  lowercase_							,  lowercase_	:				Tuple													=	weights[i], biases[i], self.out_projs[i]

                       lowercase_	:				Tuple													=	self._compute_logit(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)
                       lowercase_	:				str													=	nn.functional.log_softmax(__UpperCamelCase     ,dim=1			)

                       lowercase_	:				Union[str, Any]													=	head_logprob[:, -i] + tail_logprob_i
                       lowercase_	:				Optional[Any]													=	logprob_i

             return out


"""simple docstring"""







from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml







class        UpperCamelCase		(			yaml.SafeLoader  ):
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase			)	->      Optional[int]:


        '''simple docstring'''


        lowercase_	:				str													=	[self.constructed_objects[key_node] for key_node, _ in node.value]
        lowercase_	:				str													=	[tuple(__UpperCamelCase			) if isinstance(__UpperCamelCase     ,__UpperCamelCase			) else key for key in keys]
        lowercase_	:				List[Any]													=	Counter(__UpperCamelCase			)
        lowercase_	:				str													=	[key for key in counter if counter[key] > 1]
        if duplicate_keys:
             raise TypeError(f'''Got duplicate yaml keys: {duplicate_keys}'''			)







   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase=False			)	->      List[Any]:


        '''simple docstring'''


        lowercase_	:				Optional[int]													=	super().construct_mapping(__UpperCamelCase     ,deep=__UpperCamelCase			)
        self._check_no_duplicates_on_constructed_node(__UpperCamelCase			)
        return mapping






def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     str      ):
     lowercase_	:				Tuple													=	list(readme_content.splitlines()      )
     if full_content and full_content[0] == "---" and "---" in full_content[1:]:
          lowercase_	:				Dict													=	full_content[1:].index('---'      ) + 1
          lowercase_	:				Optional[int]													=	'\n'.join(full_content[1:sep_idx]      )
          return yamlblock, "\n".join(full_content[sep_idx + 1 :]      )

     return None, "\n".join(__SCREAMING_SNAKE_CASE      )







class        UpperCamelCase		(			lowercase_  ):
   # class attributes
   lowercase							=  {'train_eval_index'}  # train-eval-index in the YAML metadata
   @classmethod
   def 			_UpperCAmelCase   (			cls     ,__UpperCamelCase			)	->      "DatasetMetadata":


        '''simple docstring'''


        with open(__UpperCamelCase     ,encoding='utf-8'			) as readme_file:
             lowercase_							,  lowercase_	:				Optional[int]													=	_split_yaml_from_readme(readme_file.read()			)
        if yaml_string is not None:
             return cls.from_yaml_string(__UpperCamelCase			)
        else:
             return cls()
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase			)	->      List[Any]:


        '''simple docstring'''


        if path.exists():
             with open(__UpperCamelCase     ,encoding='utf-8'			) as readme_file:
                  lowercase_	:				Dict													=	readme_file.read()
        else:
             lowercase_	:				int													=	None
        lowercase_	:				Any													=	self._to_readme(__UpperCamelCase			)
        with open(__UpperCamelCase     ,'w'     ,encoding='utf-8'			) as readme_file:
             readme_file.write(__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase = None			)	->      str:


        '''simple docstring'''


        if readme_content is not None:
             lowercase_							,  lowercase_	:				Optional[Any]													=	_split_yaml_from_readme(__UpperCamelCase			)
             lowercase_	:				Optional[Any]													=	'---\n' + self.to_yaml_string() + '---\n' + content
        else:
             lowercase_	:				Tuple													=	'---\n' + self.to_yaml_string() + '---\n'
        return full_content
   @classmethod
   def 			_UpperCAmelCase   (			cls     ,__UpperCamelCase			)	->      "DatasetMetadata":


        '''simple docstring'''


        lowercase_	:				List[str]													=	yaml.load(__UpperCamelCase     ,Loader=_NoDuplicateSafeLoader			) or {}

        # Convert the YAML keys to DatasetMetadata fields
        lowercase_	:				Dict													=	{
            (key.replace('-'     ,'_'			) if key.replace('-'     ,'_'			) in cls._FIELDS_WITH_DASHES else key): value
            for key, value in metadata_dict.items()
        }
        return cls(**__UpperCamelCase			)







   def 			_UpperCAmelCase   (			self			)	->      str:


        '''simple docstring'''


        return yaml.safe_dump(
            {
                (key.replace('_'     ,'-'			) if key in self._FIELDS_WITH_DASHES else key): value
                for key, value in self.items()
            }     ,sort_keys=__UpperCamelCase     ,allow_unicode=__UpperCamelCase     ,encoding='utf-8'     ,).decode('utf-8'			)






__SCREAMING_SNAKE_CASE    	={
    "image-classification": [],
    "translation": [],
    "image-segmentation": [],
    "fill-mask": [],
    "automatic-speech-recognition": [],
    "token-classification": [],
    "sentence-similarity": [],
    "audio-classification": [],
    "question-answering": [],
    "summarization": [],
    "zero-shot-classification": [],
    "table-to-text": [],
    "feature-extraction": [],
    "other": [],
    "multiple-choice": [],
    "text-classification": [],
    "text-to-image": [],
    "text2text-generation": [],
    "zero-shot-image-classification": [],
    "tabular-classification": [],
    "tabular-regression": [],
    "image-to-image": [],
    "tabular-to-text": [],
    "unconditional-image-generation": [],
    "text-retrieval": [],
    "text-to-speech": [],
    "object-detection": [],
    "audio-to-audio": [],
    "text-generation": [],
    "conversational": [],
    "table-question-answering": [],
    "visual-question-answering": [],
    "image-to-text": [],
    "reinforcement-learning": [],
    "voice-activity-detection": [],
    "time-series-forecasting": [],
    "document-question-answering": [],
}


if __name__ == "__main__":
       from argparse import ArgumentParser

       __SCREAMING_SNAKE_CASE    	=ArgumentParser(usage="Validate the yaml metadata block of a README.md file.")
       ap.add_argument("readme_filepath")
       __SCREAMING_SNAKE_CASE    	=ap.parse_args()

       __SCREAMING_SNAKE_CASE    	=Path(args.readme_filepath)
       __SCREAMING_SNAKE_CASE    	=DatasetMetadata.from_readme(readme_filepath)
       print(dataset_metadata)
       dataset_metadata.to_readme(readme_filepath)

