




'''simple docstring'''
import tensorflow as tf

from ...tf_utils import shape_list
class   __UpperCAmelCase				(     tf.keras.layers.Layer				):



					'''simple docstring'''





					def __init__(	self     :						List[Any]      ,       _lowercase     :						Tuple      ,       _lowercase     :						List[Any]      ,       _lowercase     :						Tuple      ,       _lowercase     :						Optional[Any]      ,       _lowercase     :						List[str]=1      ,       _lowercase     :						Any=False      ,       **_lowercase     :						Optional[Any])    ->							int:
								super().__init__(**_SCREAMING_SNAKE_CASE)

								A_									= vocab_size
								A_									= d_embed
								A_									= d_proj

								A_									= cutoffs + [vocab_size]
								A_									= [0] + self.cutoffs
								A_									= div_val

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

								A_									= []
								A_									= []



					def     __snake_case						(	self     :						List[str]      ,       _lowercase     :						List[str])    ->							str:





								if self.n_clusters > 0:
											A_									= self.add_weight(
											    shape=(self.n_clusters, self.d_embed)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name='cluster_weight')
											A_									= self.add_weight(
											    shape=(self.n_clusters,)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name='cluster_bias')

								if self.div_val == 1:
											for i in range(len(self.cutoffs)):
														if self.d_proj != self.d_embed:
																	A_									= self.add_weight(
																	    shape=(self.d_embed, self.d_proj)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_projs_._{i}'      ,       )
																	self.out_projs.append(_SCREAMING_SNAKE_CASE)
														else:
																	self.out_projs.append(_SCREAMING_SNAKE_CASE)
														A_									= self.add_weight(
														    shape=(self.vocab_size, self.d_embed)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_layers_._{i}_._weight'      ,       )
														A_									= self.add_weight(
														    shape=(self.vocab_size,)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_layers_._{i}_._bias'      ,       )
														self.out_layers.append((weight, bias))
								else:
											for i in range(len(self.cutoffs)):
														A_		, A_									= self.cutoff_ends[i], self.cutoff_ends[i + 1]
														A_									= self.d_embed // (self.div_val**i)

														A_									= self.add_weight(
														    shape=(d_emb_i, self.d_proj)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_projs_._{i}')
														self.out_projs.append(_SCREAMING_SNAKE_CASE)
														A_									= self.add_weight(
														    shape=(r_idx - l_idx, d_emb_i)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_layers_._{i}_._weight'      ,       )
														A_									= self.add_weight(
														    shape=(r_idx - l_idx,)      ,       initializer='zeros'      ,       trainable=_SCREAMING_SNAKE_CASE      ,       name=F'out_layers_._{i}_._bias'      ,       )
														self.out_layers.append((weight, bias))
								super().build(_SCREAMING_SNAKE_CASE)



					@staticmethod
					def     __snake_case						(	_lowercase     :						List[str]      ,       _lowercase     :						Optional[Any]      ,       _lowercase     :						str      ,       _lowercase     :						Dict=None)    ->							Optional[Any]:
								A_									= x
								if proj is not None:
											A_									= tf.einsum('ibd,ed->ibe'      ,       _SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
								return tf.einsum('ibd,nd->ibn'      ,       _SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE) + b



					@staticmethod
					def     __snake_case						(	_lowercase     :						List[str]      ,       _lowercase     :						Optional[int])    ->							List[str]:
								A_									= shape_list(_SCREAMING_SNAKE_CASE)
								A_									= tf.range(lp_size[0]      ,       dtype=target.dtype)
								A_									= tf.stack([r, target]      ,       1)
								return tf.gather_nd(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)



					def     __snake_case						(	self     :						Any      ,       _lowercase     :						Union[str, Any]      ,       _lowercase     :						List[str]      ,       _lowercase     :						Dict=True      ,       _lowercase     :						Any=False)    ->							List[Any]:
								A_									= 0
								if self.n_clusters == 0:
											A_									= self._logit(_SCREAMING_SNAKE_CASE      ,       self.out_layers[0][0]      ,       self.out_layers[0][1]      ,       self.out_projs[0])
											if target is not None:
														A_									= tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_SCREAMING_SNAKE_CASE      ,       logits=_SCREAMING_SNAKE_CASE)
											A_									= tf.nn.log_softmax(_SCREAMING_SNAKE_CASE      ,       axis=-1)
								else:
											A_									= shape_list(_SCREAMING_SNAKE_CASE)
											A_									= []
											A_									= tf.zeros(hidden_sizes[:2])
											for i in range(len(self.cutoffs)):
														A_		, A_									= self.cutoff_ends[i], self.cutoff_ends[i + 1]
														if target is not None:
																	A_									= (target >= l_idx) & (target < r_idx)
																	A_									= tf.where(_SCREAMING_SNAKE_CASE)
																	A_									= tf.boolean_mask(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE) - l_idx

														if self.div_val == 1:
																	A_									= self.out_layers[0][0][l_idx:r_idx]
																	A_									= self.out_layers[0][1][l_idx:r_idx]
														else:
																	A_									= self.out_layers[i][0]
																	A_									= self.out_layers[i][1]

														if i == 0:
																	A_									= tf.concat([cur_W, self.cluster_weight]      ,       0)
																	A_									= tf.concat([cur_b, self.cluster_bias]      ,       0)

																	A_									= self._logit(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      ,       self.out_projs[0])
																	A_									= tf.nn.log_softmax(_SCREAMING_SNAKE_CASE)
																	out.append(head_logprob[..., : self.cutoffs[0]])
																	if target is not None:
																				A_									= tf.boolean_mask(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
																				A_									= self._gather_logprob(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
														else:
																	A_									= self._logit(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      ,       self.out_projs[i])
																	A_									= tf.nn.log_softmax(_SCREAMING_SNAKE_CASE)
																	A_									= self.cutoffs[0] + i - 1  # No probability for the head cluster
																	A_									= head_logprob[..., cluster_prob_idx, None] + tail_logprob
																	out.append(_SCREAMING_SNAKE_CASE)
																	if target is not None:
																				A_									= tf.boolean_mask(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
																				A_									= tf.boolean_mask(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
																				A_									= self._gather_logprob(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE)
																				cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
														if target is not None:
																	loss += tf.scatter_nd(_SCREAMING_SNAKE_CASE      ,       -cur_logprob      ,       shape_list(_SCREAMING_SNAKE_CASE))
											A_									= tf.concat(_SCREAMING_SNAKE_CASE      ,       axis=-1)

								if target is not None:
											if return_mean:
														A_									= tf.reduce_mean(_SCREAMING_SNAKE_CASE)
											# Add the training-time loss value to the layer using `self.add_loss()`.
											self.add_loss(_SCREAMING_SNAKE_CASE)

											# Log the loss as a metric (we could log arbitrary metrics,
											# including different metrics for training and inference.
											self.add_metric(_SCREAMING_SNAKE_CASE      ,       name=self.name      ,       aggregation='mean' if return_mean else '')

								return out








import logging
import os
from dataclasses import dataclass
from typing import List, Optional, Union

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


lowerCAmelCase__    :List[str]       =				logging.getLogger(__name__)

@dataclass(frozen=UpperCAmelCase     )
class 						__a :
		_a      :    str
		_a      :    str
		_a      :    Optional[str]           =      None
		_a      :    Optional[str]           =      None
		_a      :    Optional[str]           =      None



@dataclass(frozen=UpperCAmelCase     )
class 						__a :
		_a      :    List[int]
		_a      :    Optional[List[int]]           =      None
		_a      :    Optional[List[int]]           =      None
		_a      :    Optional[Union[int, float]]           =      None
		_a      :    Optional[int]           =      None


if is_torch_available():
		import torch
		from torch.utils.data import Dataset

		class 						__a (							UpperCAmelCase     ):
				_a      :    List[InputFeatures]
				def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE = False     ,       )					->			List[Any]:




											"""simple docstring"""





											_UpperCAmelCase											=     hans_processors[task]()

											_UpperCAmelCase											=     os.path.join(
											    _SCREAMING_SNAKE_CASE     ,       'cached_{}_{}_{}_{}'.format(
											        'dev' if evaluate else 'train'     ,       tokenizer.__class__.__name__     ,       str(_SCREAMING_SNAKE_CASE		)     ,       _SCREAMING_SNAKE_CASE     ,       )     ,       )
											_UpperCAmelCase											=     processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																		# HACK(label indices are swapped in RoBERTa pretrained model)
																		_UpperCAmelCase						,			_UpperCAmelCase											=     label_list[2], label_list[1]
											_UpperCAmelCase											=     label_list

											# Make sure only the first process in distributed training processes the dataset,
											# and the others will use the cache.
											_UpperCAmelCase											=     cached_features_file + '.lock'
											with FileLock(_SCREAMING_SNAKE_CASE		):
																		if os.path.exists(_SCREAMING_SNAKE_CASE		) and not overwrite_cache:
																									logger.info(f'''Loading features from cached file {cached_features_file}'''		)
																									_UpperCAmelCase											=     torch.load(_SCREAMING_SNAKE_CASE		)
																		else:
																									logger.info(f'''Creating features from dataset file at {data_dir}'''		)

																									_UpperCAmelCase											=     (
																									    processor.get_dev_examples(_SCREAMING_SNAKE_CASE		) if evaluate else processor.get_train_examples(_SCREAMING_SNAKE_CASE		)
																									)

																									logger.info('Training examples: %s'     ,       len(_SCREAMING_SNAKE_CASE		)		)
																									_UpperCAmelCase											=     hans_convert_examples_to_features(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																									logger.info('Saving features into cached file %s'     ,       _SCREAMING_SNAKE_CASE		)
																									torch.save(self.features     ,       _SCREAMING_SNAKE_CASE		)
				def __len__(     self		)					->			List[str]:




											"""simple docstring"""





											return len(self.features		)
				def __getitem__(     self     ,       _SCREAMING_SNAKE_CASE		)					->			InputFeatures:




											"""simple docstring"""





											return self.features[i]






				def       UpperCAmelCase__							(     self		)					->			Optional[int]:




											"""simple docstring"""





											return self.label_list


if is_tf_available():
		import tensorflow as tf

		class 						__a :
				_a      :    List[InputFeatures]
				def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = 128     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE = False     ,       )					->			Tuple:




											"""simple docstring"""





											_UpperCAmelCase											=     hans_processors[task]()
											_UpperCAmelCase											=     processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																		# HACK(label indices are swapped in RoBERTa pretrained model)
																		_UpperCAmelCase						,			_UpperCAmelCase											=     label_list[2], label_list[1]
											_UpperCAmelCase											=     label_list

											_UpperCAmelCase											=     processor.get_dev_examples(_SCREAMING_SNAKE_CASE		) if evaluate else processor.get_train_examples(_SCREAMING_SNAKE_CASE		)
											_UpperCAmelCase											=     hans_convert_examples_to_features(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

											def gen():
																		for ex_index, ex in tqdm.tqdm(enumerate(self.features		)     ,       desc='convert examples to features'		):
																									if ex_index % 10000 == 0:
																																logger.info('Writing example %d of %d' % (ex_index, len(_SCREAMING_SNAKE_CASE		))		)

																									yield (
																									    {
																									        "example_id": 0,
																									        "input_ids": ex.input_ids,
																									        "attention_mask": ex.attention_mask,
																									        "token_type_ids": ex.token_type_ids,
																									    },
																									    ex.label,
																									)

											_UpperCAmelCase											=     tf.data.Dataset.from_generator(
											    _SCREAMING_SNAKE_CASE     ,       (
											        {
											            'example_id': tf.intaa,
											            'input_ids': tf.intaa,
											            'attention_mask': tf.intaa,
											            'token_type_ids': tf.intaa,
											        },
											        tf.intaa,
											    )     ,       (
											        {
											            'example_id': tf.TensorShape([]		),
											            'input_ids': tf.TensorShape([None, None]		),
											            'attention_mask': tf.TensorShape([None, None]		),
											            'token_type_ids': tf.TensorShape([None, None]		),
											        },
											        tf.TensorShape([]		),
											    )     ,       )
				def       UpperCAmelCase__							(     self		)					->			Dict:




											"""simple docstring"""





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




											"""simple docstring"""





											return len(self.features		)
				def __getitem__(     self     ,       _SCREAMING_SNAKE_CASE		)					->			InputFeatures:




											"""simple docstring"""





											return self.features[i]






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




											"""simple docstring"""





											return self.label_list

class 						__a (							UpperCAmelCase     ):
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""simple docstring"""





									return self._create_examples(self._read_tsv(os.path.join(_SCREAMING_SNAKE_CASE     ,       'heuristics_train_set.txt'		)		)     ,       'train'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									return self._create_examples(self._read_tsv(os.path.join(_SCREAMING_SNAKE_CASE     ,       'heuristics_evaluation_set.txt'		)		)     ,       'dev'		)
		def       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									return ["contradiction", "entailment", "neutral"]






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     []
									for i, line in enumerate(_SCREAMING_SNAKE_CASE		):
																if i == 0:
																							continue
																_UpperCAmelCase											=     '%s-%s' % (set_type, line[0])
																_UpperCAmelCase											=     line[5]
																_UpperCAmelCase											=     line[6]
																_UpperCAmelCase											=     line[7][2:] if line[7].startswith('ex'		) else line[7]
																_UpperCAmelCase											=     line[0]
																examples.append(InputExample(guid=_SCREAMING_SNAKE_CASE     ,       text_a=_SCREAMING_SNAKE_CASE     ,       text_b=_SCREAMING_SNAKE_CASE     ,       label=_SCREAMING_SNAKE_CASE     ,       pairID=_SCREAMING_SNAKE_CASE		)		)
									return examples



def 			lowerCAmelCase__					(     a__:						List[InputExample]     ,    a__:						List[str]     ,    a__:						int     ,    a__:						PreTrainedTokenizer     ,    )    ->						str:


							'''simple docstring'''





							_UpperCAmelCase											=     {label: i for i, label in enumerate(a__			)}

							_UpperCAmelCase											=     []
							for ex_index, example in tqdm.tqdm(enumerate(a__			)     ,    desc='convert examples to features'			):
														if ex_index % 1_0_0_0_0 == 0:
																					logger.info('Writing example %d' % (ex_index)			)

														_UpperCAmelCase											=     tokenizer(
														    example.text_a     ,    example.text_b     ,    add_special_tokens=a__     ,    max_length=a__     ,    padding='max_length'     ,    truncation=a__     ,    return_overflowing_tokens=a__     ,    )

														_UpperCAmelCase											=     label_map[example.label] if example.label in label_map else 0

														_UpperCAmelCase											=     int(example.pairID			)

														features.append(InputFeatures(**a__     ,    label=a__     ,    pairID=a__			)			)

							for i, example in enumerate(examples[:5]			):
														logger.info('*** Example ***'			)
														logger.info(F'''guid: {example}'''			)
														logger.info(F'''features: {features[i]}'''			)

							return features


lowerCAmelCase__    :Optional[Any]       =				{
    '''hans''': 3,
}







lowerCAmelCase__    :Optional[int]       =				{
    '''hans''': HansProcessor,
}
