




'''simple docstring'''
import math
from typing import Optional

import torch
import torch.utils.checkpoint as checkpoint
from torch import nn

from ....modeling_utils import PreTrainedModel
from ....utils import add_start_docstrings, logging
from ...bert.modeling_bert import BertModel
from .configuration_retribert import RetriBertConfig


A							:				Optional[Any]			     =	logging.get_logger(__name__)

A							:				Optional[int]			     =	[
    """yjernite/retribert-base-uncased""",
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]






class 	SCREAMING_SNAKE_CASE(    __A					):
       snake_case_   = RetriBertConfig
       snake_case_   = None
       snake_case_   = """retribert"""

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Union[str, Any]:

           """simple docstring"""



           if isinstance(lowerCamelCase__ ,		nn.Linear					):
               module.weight.data.normal_(mean=0.0 ,		std=self.config.initializer_range					)
               if module.bias is not None:
                   module.bias.data.zero_()
           elif isinstance(lowerCamelCase__ ,		nn.Embedding					):
               module.weight.data.normal_(mean=0.0 ,		std=self.config.initializer_range					)
               if module.padding_idx is not None:
                   module.weight.data[module.padding_idx].zero_()
           elif isinstance(lowerCamelCase__ ,		nn.LayerNorm					):
               module.bias.data.zero_()
               module.weight.data.fill_(1.0					)

A							:				List[Any]			     =	r"""

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""






@add_start_docstrings(
    """Bert Based model to embed queries or document for document retrieval."""       ,				__A       ,				)
class 	SCREAMING_SNAKE_CASE(    __A					):

       def __init__(   self ,		lowerCamelCase__					)   -> None:

           """simple docstring"""



           super().__init__(lowerCamelCase__					)
           __lowercase          =						config.projection_dim

           __lowercase          =						BertModel(lowerCamelCase__					)
           __lowercase          =						None if config.share_encoders else BertModel(lowerCamelCase__					)
           __lowercase          =						nn.Dropout(config.hidden_dropout_prob					)
           __lowercase          =						nn.Linear(config.hidden_size ,		config.projection_dim ,		bias=lowerCamelCase__					)
           __lowercase          =						nn.Linear(config.hidden_size ,		config.projection_dim ,		bias=lowerCamelCase__					)

           __lowercase          =						nn.CrossEntropyLoss(reduction="""mean"""					)

           # Initialize weights and apply final processing
           self.post_init()

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__=-1 ,		)   -> Optional[int]:

           """simple docstring"""



           if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
               return sent_encoder(lowerCamelCase__ ,		attention_mask=lowerCamelCase__					)[1]
           else:
               # prepare implicit variables
               __lowercase          =						input_ids.device
               __lowercase          =						input_ids.size()
               __lowercase          =						torch.zeros(lowerCamelCase__ ,		dtype=torch.long ,		device=lowerCamelCase__					)
               __lowercase          =						[None] * sent_encoder.config.num_hidden_layers
               __lowercase          =						sent_encoder.get_extended_attention_mask(
                   lowerCamelCase__ ,		lowerCamelCase__					)

               # define function for checkpointing
               def partial_encode(*lowerCamelCase__					):
                   __lowercase          =						sent_encoder.encoder(
                       inputs[0] ,		attention_mask=inputs[1] ,		head_mask=lowerCamelCase__ ,		)
                   __lowercase          =						encoder_outputs[0]
                   __lowercase          =						sent_encoder.pooler(lowerCamelCase__					)
                   return pooled_output

               # run embedding layer on everything at once
               __lowercase          =						sent_encoder.embeddings(
                   input_ids=lowerCamelCase__ ,		position_ids=lowerCamelCase__ ,		token_type_ids=lowerCamelCase__ ,		inputs_embeds=lowerCamelCase__					)
               # run encoding and pooling on one mini-batch at a time
               __lowercase          =						[]
               for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size					)					):
                   __lowercase          =						embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                   __lowercase          =						extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                   __lowercase          =						checkpoint.checkpoint(lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)
                   pooled_output_list.append(lowerCamelCase__					)
               return torch.cat(lowerCamelCase__ ,		dim=0					)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=None ,		lowerCamelCase__=-1 ,		)   -> Dict:

           """simple docstring"""



           __lowercase          =						self.embed_sentences_checkpointed(
               lowerCamelCase__ ,		lowerCamelCase__ ,		self.bert_query ,		lowerCamelCase__ ,		)
           return self.project_query(lowerCamelCase__					)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=None ,		lowerCamelCase__=-1 ,		)   -> Optional[Any]:

           """simple docstring"""



           __lowercase          =						self.embed_sentences_checkpointed(
               lowerCamelCase__ ,		lowerCamelCase__ ,		self.bert_query if self.bert_doc is None else self.bert_doc ,		lowerCamelCase__ ,		)
           return self.project_doc(lowerCamelCase__					)




       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ = -1 ,		)   -> torch.FloatTensor:

           """simple docstring"""



           __lowercase          =						input_ids_query.device
           __lowercase          =						self.embed_questions(lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)
           __lowercase          =						self.embed_answers(lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)
           __lowercase          =						torch.mm(lowerCamelCase__ ,		a_reps.t()					)
           __lowercase          =						self.ce_loss(lowerCamelCase__ ,		torch.arange(compare_scores.shape[1]					).to(lowerCamelCase__					)					)
           __lowercase          =						self.ce_loss(compare_scores.t() ,		torch.arange(compare_scores.shape[0]					).to(lowerCamelCase__					)					)
           __lowercase          =						(loss_qa + loss_aq) / 2
           return loss










'''simple docstring'''
import json
import os
import re
import unicodedata
from typing import Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
from ...utils import logging


A							:				int			     =	logging.get_logger(__name__)

A							:				List[Any]			     =	{
    """vocab_file""": """vocab.json""",
    """merges_file""": """merges.txt""",
}

A							:				Tuple			     =	{
    """vocab_file""": {"""openai-gpt""": """https://huggingface.co/openai-gpt/resolve/main/vocab.json"""},
    """merges_file""": {"""openai-gpt""": """https://huggingface.co/openai-gpt/resolve/main/merges.txt"""},
}

A							:				Any			     =	{
    """openai-gpt""": 512,
}
def     snake_case_  (       a__		:     List[str]      ):



    """simple docstring"""




    __lowercase          =						text.strip()
    if not text:
        return []
    __lowercase          =						text.split()
    return tokens






class 	SCREAMING_SNAKE_CASE(    __A					):

       def __init__(   self ,		lowerCamelCase__=True ,		lowerCamelCase__=None ,		lowerCamelCase__=True ,		lowerCamelCase__=None ,		lowerCamelCase__=True ,		)   -> Optional[Any]:

           """simple docstring"""



           if never_split is None:
               __lowercase          =						[]
           __lowercase          =						do_lower_case
           __lowercase          =						set(lowerCamelCase__					)
           __lowercase          =						tokenize_chinese_chars
           __lowercase          =						strip_accents
           __lowercase          =						do_split_on_punc

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=None					)   -> Optional[Any]:

           """simple docstring"""



           __lowercase          =						self.never_split.union(set(lowerCamelCase__					)					) if never_split else self.never_split
           __lowercase          =						self._clean_text(lowerCamelCase__					)

           # This was added on November 1st, 2018 for the multilingual and Chinese
           # models. This is also applied to the English models now, but it doesn't
           # matter since the English models were not trained on any Chinese data
           # and generally don't have any Chinese data in them (there are Chinese
           # characters in the vocabulary because Wikipedia does have some Chinese
           # words in the English Wikipedia.).
           if self.tokenize_chinese_chars:
               __lowercase          =						self._tokenize_chinese_chars(lowerCamelCase__					)
           # prevents treating the same character with different unicode codepoints as different characters
           __lowercase          =						unicodedata.normalize("""NFC""" ,		lowerCamelCase__					)
           __lowercase          =						whitespace_tokenize(lowerCamelCase__					)
           __lowercase          =						[]
           for token in orig_tokens:
               if token not in never_split:
                   if self.do_lower_case:
                       __lowercase          =						token.lower()
                       if self.strip_accents is not False:
                           __lowercase          =						self._run_strip_accents(lowerCamelCase__					)
                   elif self.strip_accents:
                       __lowercase          =						self._run_strip_accents(lowerCamelCase__					)
               split_tokens.extend(self._run_split_on_punc(lowerCamelCase__ ,		lowerCamelCase__					)					)

           __lowercase          =						whitespace_tokenize(""" """.join(lowerCamelCase__					)					)
           return output_tokens

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Any:

           """simple docstring"""



           __lowercase          =						unicodedata.normalize("""NFD""" ,		lowerCamelCase__					)
           __lowercase          =						[]
           for char in text:
               __lowercase          =						unicodedata.category(lowerCamelCase__					)
               if cat == "Mn":
                   continue
               output.append(lowerCamelCase__					)
           return "".join(lowerCamelCase__					)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=None					)   -> Optional[int]:

           """simple docstring"""



           if not self.do_split_on_punc or (never_split is not None and text in never_split):
               return [text]
           __lowercase          =						list(lowerCamelCase__					)
           __lowercase          =						0
           __lowercase          =						True
           __lowercase          =						[]
           while i < len(lowerCamelCase__					):
               __lowercase          =						chars[i]
               if _is_punctuation(lowerCamelCase__					):
                   output.append([char]					)
                   __lowercase          =						True
               else:
                   if start_new_word:
                       output.append([]					)
                   __lowercase          =						False
                   output[-1].append(lowerCamelCase__					)
               i += 1

           return ["".join(lowerCamelCase__					) for x in output]

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Any:

           """simple docstring"""



           __lowercase          =						[]
           for char in text:
               __lowercase          =						ord(lowerCamelCase__					)
               if self._is_chinese_char(lowerCamelCase__					):
                   output.append(""" """					)
                   output.append(lowerCamelCase__					)
                   output.append(""" """					)
               else:
                   output.append(lowerCamelCase__					)
           return "".join(lowerCamelCase__					)

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Tuple:

           """simple docstring"""



           if (
               (cp >= 0X4_e00 and cp <= 0X9_fff)
               or (cp >= 0X3_400 and cp <= 0X4_dbf)  #
               or (cp >= 0X20_000 and cp <= 0X2a_6df)  #
               or (cp >= 0X2a_700 and cp <= 0X2b_73f)  #
               or (cp >= 0X2b_740 and cp <= 0X2b_81f)  #
               or (cp >= 0X2b_820 and cp <= 0X2c_eaf)  #
               or (cp >= 0Xf_900 and cp <= 0Xf_aff)
               or (cp >= 0X2f_800 and cp <= 0X2f_a1f)  #
           ):  #
               return True

           return False




       def    snake_case__      (   self ,		lowerCamelCase__					)   -> str:

           """simple docstring"""



           __lowercase          =						[]
           for char in text:
               __lowercase          =						ord(lowerCamelCase__					)
               if cp == 0 or cp == 0Xf_ffd or _is_control(lowerCamelCase__					):
                   continue
               if _is_whitespace(lowerCamelCase__					):
                   output.append(""" """					)
               else:
                   output.append(lowerCamelCase__					)
           return "".join(lowerCamelCase__					)

def     snake_case_  (       a__		:     List[Any]      ):



    """simple docstring"""




    __lowercase          =						set()
    __lowercase          =						word[0]
    for char in word[1:]:
        pairs.add((prev_char, char)      )
        __lowercase          =						char
    return pairs
def     snake_case_  (       a__		:     Union[str, Any]      ):



    """simple docstring"""




    __lowercase          =						text.replace("""—"""   ,"""-"""      )
    __lowercase          =						text.replace("""–"""   ,"""-"""      )
    __lowercase          =						text.replace("""―"""   ,"""-"""      )
    __lowercase          =						text.replace("""…"""   ,"""..."""      )
    __lowercase          =						text.replace("""´"""   ,"""'"""      )
    __lowercase          =						re.sub(r"""(-+|~+|!+|\"+|;+|\?+|\++|,+|\)+|\(+|\\+|\/+|\*+|\[+|\]+|}+|{+|\|+|_+)"""   ,r""" \1 """   ,a__      )
    __lowercase          =						re.sub(r"""\s*\n\s*"""   ,""" \n """   ,a__      )
    __lowercase          =						re.sub(r"""[^\S\n]+"""   ,""" """   ,a__      )
    return text.strip()






class 	SCREAMING_SNAKE_CASE(    __A					):
       snake_case_			:					Dict   = VOCAB_FILES_NAMES
       snake_case_			:					Tuple   = PRETRAINED_VOCAB_FILES_MAP
       snake_case_			:					List[Any]   = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       snake_case_			:					Union[str, Any]   = ["""input_ids""", """attention_mask"""]

       def __init__(   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__="<unk>" ,		**lowerCamelCase__					)   -> List[Any]:

           """simple docstring"""



           super().__init__(unk_token=lowerCamelCase__ ,		**lowerCamelCase__					)

           try:
               import ftfy
               from spacy.lang.en import English

               __lowercase          =						English()
               __lowercase          =						_nlp.tokenizer
               __lowercase          =						ftfy.fix_text
           except ImportError:
               logger.warning("""ftfy or spacy is not installed using BERT BasicTokenizer instead of SpaCy & ftfy."""					)
               __lowercase          =						BasicTokenizer(do_lower_case=lowerCamelCase__					)
               __lowercase          =						None

           with open(lowerCamelCase__ ,		encoding="""utf-8"""					) as vocab_handle:
               __lowercase          =						json.load(lowerCamelCase__					)
           __lowercase          =						{v: k for k, v in self.encoder.items()}
           with open(lowerCamelCase__ ,		encoding="""utf-8"""					) as merges_handle:
               __lowercase          =						merges_handle.read().split("""\n"""					)[1:-1]
           __lowercase          =						[tuple(merge.split()					) for merge in merges]
           __lowercase          =						dict(zip(lowerCamelCase__ ,		range(len(lowerCamelCase__					)					)					)					)
           __lowercase          =						{}

       @property
       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           return True

       @property
       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           return len(self.encoder					)

       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           return dict(self.encoder ,		**self.added_tokens_encoder					)

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> int:

           """simple docstring"""



           __lowercase          =						tuple(token[:-1]					) + (token[-1] + """</w>""",)
           if token in self.cache:
               return self.cache[token]
           __lowercase          =						get_pairs(lowerCamelCase__					)

           if not pairs:
               return token + "</w>"

           while True:
               __lowercase          =						min(lowerCamelCase__ ,		key=lambda lowerCamelCase__					: self.bpe_ranks.get(lowerCamelCase__ ,		float("""inf"""					)					)					)
               if bigram not in self.bpe_ranks:
                   break
               __lowercase    ,__lowercase          =						bigram
               __lowercase          =						[]
               __lowercase          =						0
               while i < len(lowerCamelCase__					):
                   try:
                       __lowercase          =						word.index(lowerCamelCase__ ,		lowerCamelCase__					)
                   except ValueError:
                       new_word.extend(word[i:]					)
                       break
                   else:
                       new_word.extend(word[i:j]					)
                       __lowercase          =						j

                   if word[i] == first and i < len(lowerCamelCase__					) - 1 and word[i + 1] == second:
                       new_word.append(first + second					)
                       i += 2
                   else:
                       new_word.append(word[i]					)
                       i += 1
               __lowercase          =						tuple(lowerCamelCase__					)
               __lowercase          =						new_word
               if len(lowerCamelCase__					) == 1:
                   break
               else:
                   __lowercase          =						get_pairs(lowerCamelCase__					)
           __lowercase          =						""" """.join(lowerCamelCase__					)
           if word == "\n  </w>":
               __lowercase          =						"""\n</w>"""
           __lowercase          =						word
           return word

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> List[str]:

           """simple docstring"""



           __lowercase          =						[]
           if self.fix_text is None:
               # Using BERT's BasicTokenizer
               __lowercase          =						self.nlp.tokenize(lowerCamelCase__					)
               for token in text:
                   split_tokens.extend(list(self.bpe(lowerCamelCase__					).split(""" """					)					)					)
           else:
               # Using SpaCy & ftfy (original tokenization process of OpenAI GPT)
               __lowercase          =						self.nlp(text_standardize(self.fix_text(lowerCamelCase__					)					)					)
               for token in text:
                   split_tokens.extend(list(self.bpe(token.text.lower()					).split(""" """					)					)					)
           return split_tokens

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> List[str]:

           """simple docstring"""



           return self.encoder.get(lowerCamelCase__ ,		self.encoder.get(self.unk_token					)					)

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> str:

           """simple docstring"""



           return self.decoder.get(lowerCamelCase__ ,		self.unk_token					)

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Any:

           """simple docstring"""



           __lowercase          =						"""""".join(lowerCamelCase__					).replace("""</w>""" ,		""" """					).strip()
           return out_string




       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ = None					)   -> Tuple[str]:

           """simple docstring"""



           if not os.path.isdir(lowerCamelCase__					):
               logger.error(F'Vocabulary path ({save_directory}) should be a directory'					)
               return
           __lowercase          =						os.path.join(
               lowerCamelCase__ ,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]					)
           __lowercase          =						os.path.join(
               lowerCamelCase__ ,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]					)

           with open(lowerCamelCase__ ,		"""w""" ,		encoding="""utf-8"""					) as f:
               f.write(json.dumps(self.encoder ,		indent=2 ,		sort_keys=lowerCamelCase__ ,		ensure_ascii=lowerCamelCase__					) + """\n"""					)

           __lowercase          =						0
           with open(lowerCamelCase__ ,		"""w""" ,		encoding="""utf-8"""					) as writer:
               writer.write("""#version: 0.2\n"""					)
               for bpe_tokens, token_index in sorted(self.bpe_ranks.items() ,		key=lambda lowerCamelCase__					: kv[1]					):
                   if index != token_index:
                       logger.warning(
                           F'Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'
                           """ Please check that the tokenizer is not corrupted!"""					)
                       __lowercase          =						token_index
                   writer.write(""" """.join(lowerCamelCase__					) + """\n"""					)
                   index += 1

           return vocab_file, merge_file




