


'''simple docstring'''






import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


__A	:	int     =     {
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warning': logging.WARNING,
    'error': logging.ERROR,
    'critical': logging.CRITICAL,
}

__A	:	List[str]     =     logging.WARNING





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							Optional[int]			       =						os.getenv("""DATASETS_VERBOSITY"""				,      lowerCamelCase_ )
		if env_level_str:
				if env_level_str in log_levels:
						return log_levels[env_level_str]
				else:
						logging.getLogger().warning(
						    F'''Unknown option DATASETS_VERBOSITY={env_level_str}, '''
						    F'''has to be one of: { ", ".join(log_levels.keys() ) }''' )
		return _default_log_level





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return __name__.split(""".""" )[0]





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return logging.getLogger(_get_library_name() )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		# Apply our default configuration to the library root logger.
		snake_case_		:							int			       =						_get_library_root_logger()
		library_root_logger.setLevel(_get_default_logging_level() )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							str			       =						_get_library_root_logger()
		library_root_logger.setLevel(logging.NOTSET )





def      UpperCAmelCase      (		lowerCamelCase_  :Optional[str] = None ):







		'''simple docstring'''
		if name is None:
				snake_case_		:							List[str]			       =						_get_library_name()
		return logging.getLogger(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return _get_library_root_logger().getEffectiveLevel()





def      UpperCAmelCase      (		lowerCamelCase_  :int ):







		'''simple docstring'''
		_get_library_root_logger().setLevel(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return set_verbosity(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return set_verbosity(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return set_verbosity(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		return set_verbosity(lowerCamelCase_ )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							Optional[int]			       =						False





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							Any			       =						True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()
class 						__UpperCamelCase	:





							def __init__(  self		:Optional[Any]  ,*_UpperCamelCase		:List[Any]  ,**_UpperCamelCase		:List[Any]			):  # pylint: disable=unused-argument
									snake_case_		:							Optional[int]			       =						args[0] if args else None





							def __iter__(  self		:List[Any]			):
									return iter(self._iterator			)





							def __getattr__(  self		:int  ,_UpperCamelCase		:int			):

									def empty_fn(*_UpperCamelCase		:Optional[int]  ,**_UpperCamelCase		:int			):  # pylint: disable=unused-argument
											return

									return empty_fn





							def __enter__(  self		:Tuple			):
									return self





							def __exit__(  self		:Optional[Any]  ,_UpperCamelCase		:Optional[int]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Optional[Any]			):
									return





__A	:	List[Any]     =     True
class 						__UpperCamelCase	:





							def __call__(  self		:List[Any]  ,*_UpperCamelCase		:int  ,_UpperCamelCase		:Any=False  ,**_UpperCamelCase		:Optional[Any]			):
									if _tqdm_active and not disable:
											return tqdm_lib.tqdm(*_UpperCamelCase  ,**_UpperCamelCase			)
									else:
											return EmptyTqdm(*_UpperCamelCase  ,**_UpperCamelCase			)





							def 		a__   (  self		:Any  ,*_UpperCamelCase		:Optional[int]  ,**_UpperCamelCase		:Tuple			):
									snake_case_		:							int			       =						None
									if _tqdm_active:
											return tqdm_lib.tqdm.set_lock(*_UpperCamelCase  ,**_UpperCamelCase			)





							def 		a__   (  self		:Tuple			):
									if _tqdm_active:
											return tqdm_lib.tqdm.get_lock()





__A	:	str     =     _tqdm_cls()





def      UpperCAmelCase      (		):







		'''simple docstring'''
		global _tqdm_active
		return bool(_tqdm_active )





def      UpperCAmelCase      (		):







		'''simple docstring'''
		global _tqdm_active
		snake_case_		:							Tuple			       =						True






def      UpperCAmelCase      (		):







		'''simple docstring'''
		global _tqdm_active
		snake_case_		:							Union[str, Any]			       =						False



'''simple docstring'''






import torch

from transformers import AutoModel
class 						__UpperCamelCase	(						torch.nn.Module							):





							def __init__(  self		:Union[str, Any]  ,_UpperCamelCase		:Tuple="sayef/fsner-bert-base-uncased"			):
									super(_UpperCamelCase  ,self			).__init__()

									snake_case_		:							Dict			       =						AutoModel.from_pretrained(_UpperCamelCase  ,return_dict=_UpperCamelCase			)
									snake_case_		:							Dict			       =						torch.nn.CosineSimilarity(3  ,1E-0_8			)
									snake_case_		:							int			       =						torch.nn.Softmax(dim=1			)





							def 		a__   (  self		:Optional[Any]  ,**_UpperCamelCase		:Optional[Any]			):
									return self.bert(**_UpperCamelCase			).last_hidden_state





							def 		a__   (  self		:List[Any]  ,_UpperCamelCase		:Dict			):
									return token_embeddings.sum(2  ,keepdim=_UpperCamelCase			)





							def 		a__   (  self		:Optional[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:List[str]=1			):
									return self.softmax(T * self.cos(_UpperCamelCase  ,_UpperCamelCase			)			)





							def 		a__   (  self		:Optional[Any]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:int			):
									snake_case_		:							Union[str, Any]			       =						W_supports["""sizes"""].tolist()
									snake_case_		:							Dict			       =						W_supports["""start_token_id"""].item()
									snake_case_		:							int			       =						W_supports["""end_token_id"""].item()

									del W_supports["sizes"]
									del W_supports["start_token_id"]
									del W_supports["end_token_id"]

									snake_case_		:							int			       =						self.BERT(**_UpperCamelCase			)
									snake_case_		:							int			       =						self.BERT(**_UpperCamelCase			)

									snake_case_		:							Any			       =						None
									snake_case_		:							Any			       =						None

									snake_case_		:							Union[str, Any]			       =						W_supports["""input_ids"""] == start_token_id
									snake_case_		:							Tuple			       =						W_supports["""input_ids"""] == end_token_id

									for i, size in enumerate(_UpperCamelCase			):
											if i == 0:
													snake_case_		:							str			       =						0
											else:
													snake_case_		:							Optional[Any]			       =						support_sizes[i - 1]

											snake_case_		:							Dict			       =						S[s : s + size][start_token_masks[s : s + size]]
											snake_case_		:							Optional[Any]			       =						S[s : s + size][end_token_masks[s : s + size]]

											snake_case_		:							int			       =						torch.matmul(q[i]  ,s_start.T			).sum(1			).softmax(0			)
											snake_case_		:							List[str]			       =						torch.matmul(q[i]  ,s_end.T			).sum(1			).softmax(0			)

											if p_starts is not None:
													snake_case_		:							List[str]			       =						torch.vstack((p_starts, p_start)			)
													snake_case_		:							Tuple			       =						torch.vstack((p_ends, p_end)			)
											else:
													snake_case_		:							Any			       =						p_start
													snake_case_		:							List[str]			       =						p_end

									return p_starts, p_ends