

'''simple docstring'''

import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase							=					logging.get_logger(__name__)

UpperCAmelCase							=					[
    (
        "protein1",
        "MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA",
    ),
    ("protein2", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"),
    ("protein3", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"),
    ("protein4", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"),
]

UpperCAmelCase							=					{
    "esm1b_t33_650M_UR50S": esm_module.pretrained.esm1b_t33_650M_UR50S,
    "esm1v_t33_650M_UR90S_1": esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    "esm1v_t33_650M_UR90S_2": esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    "esm1v_t33_650M_UR90S_3": esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    "esm1v_t33_650M_UR90S_4": esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    "esm1v_t33_650M_UR90S_5": esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    "esm2_t48_15B_UR50D": esm_module.pretrained.esm2_t48_15B_UR50D,
    "esm2_t36_3B_UR50D": esm_module.pretrained.esm2_t36_3B_UR50D,
    "esm2_t33_650M_UR50D": esm_module.pretrained.esm2_t33_650M_UR50D,
    "esm2_t30_150M_UR50D": esm_module.pretrained.esm2_t30_150M_UR50D,
    "esm2_t12_35M_UR50D": esm_module.pretrained.esm2_t12_35M_UR50D,
    "esm2_t6_8M_UR50D": esm_module.pretrained.esm2_t6_8M_UR50D,
    "esmfold_v1": esmfold_va,
}

UpperCAmelCase							=					list("ARNDCQEGHILKMFPSTWYV")

UpperCAmelCase							=					restypes + ["X"]
UpperCAmelCase							=					restypes_with_x + ["<pad>", "<mask>", "<cls>", "<sep>", "<eos>"]

def 			__UpperCamelCase							(		):






     '''simple docstring'''





     with TemporaryDirectory() as tempdir:
          __lowercase     						='\n'.join(lowercase__	)
          __lowercase     						=Path(lowercase__	) / 'vocab.txt'
          vocab_file.write_text(lowercase__	)
          __lowercase     						=EsmTokenizer(vocab_file=str(lowercase__	)	)
     __lowercase     						=0  # Overlaps with 'A' but that seems to be what they want
     return hf_tokenizer

def 			__UpperCamelCase							(		lowercase__  :					str,   lowercase__  :					Optional[int]	):






     '''simple docstring'''





     __lowercase     						=our_module.load_state_dict(original_module.state_dict()	)
     if status.missing_keys:
          raise ValueError(F'''Missing keys: {status.missing_keys}'''	)
     if status.unexpected_keys:
          raise ValueError(F'''Unexpected keys: {status.unexpected_keys}'''	)

def 			__UpperCamelCase							(		lowercase__  :					str,   lowercase__  :					str,   lowercase__  :					bool,   lowercase__  :					str,   lowercase__  :					str	):






     '''simple docstring'''





     if model.startswith('esmfold'	):
          __lowercase     						=MODEL_MAPPING[model]()
     else:
          __lowercase						,       __lowercase     						=MODEL_MAPPING[model]()
     esm.eval()  # disable dropout

     if model.startswith('esmfold'	):
          __lowercase     						=esm.esm.embed_dim
          __lowercase     						=esm.esm.num_layers
          __lowercase     						=esm.esm.attention_heads
          __lowercase     						=4 * embed_dim
          __lowercase     						=esm.esm.token_dropout
          __lowercase     						=False  # This code path does not exist in ESM-2
          __lowercase     						='rotary'
          __lowercase     						=True
          __lowercase     						=EsmFoldConfig()
          for key, val in esm.cfg.items():
               if hasattr(lowercase__,   lowercase__	) and key != "trunk":
                    setattr(lowercase__,   lowercase__,   lowercase__	)
          for key, val in esm.cfg.trunk.items():
               if hasattr(esmfold_config.trunk,   lowercase__	) and key != "structure_module":
                    setattr(esmfold_config.trunk,   lowercase__,   lowercase__	)
          for key, val in esm.cfg.trunk.structure_module.items():
               if hasattr(esmfold_config.trunk.structure_module,   lowercase__	):
                    setattr(esmfold_config.trunk.structure_module,   lowercase__,   lowercase__	)
     elif hasattr(lowercase__,   'args'	):
          # Indicates an ESM-1b or ESM-1v model
          __lowercase     						=esm.args.embed_dim
          __lowercase     						=esm.args.layers
          __lowercase     						=esm.args.attention_heads
          __lowercase     						=esm.args.ffn_embed_dim
          __lowercase     						=esm.args.token_dropout
          __lowercase     						=True if esm.emb_layer_norm_before else False
          __lowercase     						='absolute'
          __lowercase     						=False
          __lowercase     						=None
     else:
          # Indicates an ESM-2 model
          __lowercase     						=esm.embed_dim
          __lowercase     						=esm.num_layers
          __lowercase     						=esm.attention_heads
          __lowercase     						=4 * embed_dim  # This is hardcoded in ESM-2
          __lowercase     						=esm.token_dropout
          __lowercase     						=False  # This code path does not exist in ESM-2
          __lowercase     						='rotary'
          __lowercase     						=False
          __lowercase     						=None

     if is_folding_model:
          __lowercase     						=esm.esm.alphabet
     __lowercase     						=tuple(alphabet.all_toks	)
     __lowercase     						=alphabet.mask_idx
     __lowercase     						=alphabet.padding_idx

     if is_folding_model:
          __lowercase     						=esm.esm
     else:
          __lowercase     						=esm

     __lowercase     						=EsmConfig(
         vocab_size=original_esm_model.embed_tokens.num_embeddings,   mask_token_id=lowercase__,   hidden_size=lowercase__,   num_hidden_layers=lowercase__,   num_attention_heads=lowercase__,   intermediate_size=lowercase__,   max_position_embeddings=10_26,   layer_norm_eps=1E-5,   attention_probs_dropout_prob=0.0,   hidden_dropout_prob=0.0,   pad_token_id=lowercase__,   emb_layer_norm_before=lowercase__,   token_dropout=lowercase__,   position_embedding_type=lowercase__,   is_folding_model=lowercase__,   esmfold_config=lowercase__,   vocab_list=lowercase__,   )
     if classification_head:
          __lowercase     						=esm.classification_heads['mnli'].out_proj.weight.shape[0]
     print('Our ESM config:',   lowercase__	)

     if model.startswith('esmfold'	):
          __lowercase     						=EsmForProteinFolding
     elif classification_head:
          __lowercase     						=EsmForSequenceClassification
     else:
          __lowercase     						=EsmForMaskedLM
     __lowercase     						=model_class(lowercase__	)
     model.eval()

     # Now let's copy all the weights.
     # Embeddings
     __lowercase     						=original_esm_model.embed_tokens.weight
     if position_embedding_type == "absolute":
          __lowercase     						=original_esm_model.embed_positions.weight

     if config.emb_layer_norm_before:
          __lowercase     						=original_esm_model.emb_layer_norm_before.weight
          __lowercase     						=original_esm_model.emb_layer_norm_before.bias

     __lowercase     						=original_esm_model.emb_layer_norm_after.weight
     __lowercase     						=original_esm_model.emb_layer_norm_after.bias

     for i in range(config.num_hidden_layers	):
          # Encoder: start of layer
          __lowercase     						=model.esm.encoder.layer[i]
          # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
          __lowercase     						=original_esm_model.layers[i]

          # self attention
          __lowercase     						=layer.attention.self
          assert (
              esm_layer.self_attn.k_proj.weight.data.shape
              == esm_layer.self_attn.q_proj.weight.data.shape
              == esm_layer.self_attn.v_proj.weight.data.shape
              == torch.Size((config.hidden_size, config.hidden_size)	)
          )

          __lowercase     						=esm_layer.self_attn.q_proj.weight
          __lowercase     						=esm_layer.self_attn.q_proj.bias
          __lowercase     						=esm_layer.self_attn.k_proj.weight
          __lowercase     						=esm_layer.self_attn.k_proj.bias
          __lowercase     						=esm_layer.self_attn.v_proj.weight
          __lowercase     						=esm_layer.self_attn.v_proj.bias

          if getattr(esm_layer.self_attn,   'rot_emb',   lowercase__	) is not None:
               # Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
               # During the training of ESM-2 the model was converted to float16 precision, which also converts
               # the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
               # If we recompute inv_freq without this loss of precision then we will get subtly different rotary
               # embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
               # we make sure the new model copies the data from the old inv_freq.
               __lowercase     						=esm_layer.self_attn.rot_emb.inv_freq

          # LayerNorm changes for pre-activation
          __lowercase     						=esm_layer.self_attn_layer_norm.weight
          __lowercase     						=esm_layer.self_attn_layer_norm.bias
          __lowercase     						=esm_layer.final_layer_norm.weight
          __lowercase     						=esm_layer.final_layer_norm.bias

          # self-attention output
          __lowercase     						=layer.attention.output
          assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
          __lowercase     						=esm_layer.self_attn.out_proj.weight
          __lowercase     						=esm_layer.self_attn.out_proj.bias

          # intermediate
          __lowercase     						=layer.intermediate
          assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
          __lowercase     						=esm_layer.fca.weight
          __lowercase     						=esm_layer.fca.bias

          # output
          __lowercase     						=layer.output
          assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
          __lowercase     						=esm_layer.fca.weight
          __lowercase     						=esm_layer.fca.bias
          # end of layer

     if is_folding_model:
          __lowercase     						=esm.esm_s_combine.data
          __lowercase     						=esm.afa_to_esm.data
          transfer_and_check_weights(esm.embedding,   model.embedding	)
          transfer_and_check_weights(esm.esm_s_mlp,   model.esm_s_mlp	)
          transfer_and_check_weights(esm.trunk,   model.trunk	)
          transfer_and_check_weights(esm.distogram_head,   model.distogram_head	)
          transfer_and_check_weights(esm.ptm_head,   model.ptm_head	)
          transfer_and_check_weights(esm.lm_head,   model.lm_head	)
          transfer_and_check_weights(esm.lddt_head,   model.lddt_head	)

     elif classification_head:
          __lowercase     						=esm.esm.classification_heads['mnli'].dense.weight
          __lowercase     						=esm.classification_heads['mnli'].dense.bias
          __lowercase     						=esm.classification_heads['mnli'].out_proj.weight
          __lowercase     						=esm.classification_heads['mnli'].out_proj.bias
     else:
          # LM Head
          __lowercase     						=esm.lm_head.dense.weight
          __lowercase     						=esm.lm_head.dense.bias
          __lowercase     						=esm.lm_head.layer_norm.weight
          __lowercase     						=esm.lm_head.layer_norm.bias
          __lowercase     						=esm.lm_head.weight
          __lowercase     						=esm.lm_head.bias

     # Contact prediction head
     transfer_and_check_weights(esm.contact_head,   model.esm.contact_head	)

     # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
     if is_folding_model:
          # Folding models aren't trained on masked inputs and don't like mask tokens.
          __lowercase     						=SAMPLE_DATA[:2]
     else:
          __lowercase     						=SAMPLE_DATA

     if is_folding_model:
          __lowercase     						=get_esmfold_tokenizer()
          __lowercase     						=hf_tokenizer(
              [row[1] for row in sample_data],   return_tensors='pt',   padding=lowercase__,   add_special_tokens=lowercase__	)
          __lowercase						,       __lowercase						,       __lowercase						,       __lowercase						,       __lowercase     						=esmfold_encode_sequences([row[1] for row in sample_data]	)
          __lowercase     						=torch.all(hf_tokens['input_ids'] == esmfold_aas	) and torch.all(
              hf_tokens['attention_mask'] == esmfold_mask	)
     else:
          # Let's check that we get the same results.
          __lowercase     						=alphabet.get_batch_converter()
          __lowercase						,       __lowercase						,       __lowercase     						=batch_converter(lowercase__	)
          # Prepare tokenizer and make sure it matches
          with TemporaryDirectory() as tempdir:
               __lowercase     						='\n'.join(alphabet.all_toks	)
               __lowercase     						=Path(lowercase__	) / 'vocab.txt'
               vocab_file.write_text(lowercase__	)
               __lowercase     						=EsmTokenizer(vocab_file=str(lowercase__	)	)

          __lowercase     						=hf_tokenizer([row[1] for row in sample_data],   return_tensors='pt',   padding=lowercase__	)
          __lowercase     						=torch.all(hf_tokens['input_ids'] == batch_tokens	)

     print('Do both models tokenizers output the same tokens?',   '🔥' if success else '💩'	)
     if not success:
          raise Exception('Tokenization does not match!'	)

     with torch.no_grad():
          if is_folding_model:
               # Let's test the model in parts
               # ESMFold always converts the ESM stem to float16, which requires float16 ops
               # that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
               # ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
               # original and the converted model on the GPU at the same time.
               __lowercase     						=esm.cuda().infer([row[1] for row in sample_data]	)
               __lowercase     						=model.cuda()(
                   input_ids=hf_tokens['input_ids'].cuda(),   attention_mask=hf_tokens['attention_mask'].cuda()	)
          else:
               __lowercase     						=model(**lowercase__,   output_hidden_states=lowercase__	)
               __lowercase     						=our_output['logits']
               if classification_head:
                    __lowercase     						=esm.model.classification_heads['mnli'](esm.extract_features(lowercase__	)	)
               else:
                    __lowercase     						=esm(hf_tokens['input_ids'],   repr_layers=list(range(9_99	)	)	)
                    __lowercase     						=their_output['logits']

          if is_folding_model:
               __lowercase     						=torch.max(torch.abs(our_output['positions'] - their_output['positions']	)	).item()
               __lowercase     						=torch.allclose(our_output['positions'],   their_output['positions'],   atol=1E-5	)
          else:
               __lowercase     						=torch.max(torch.abs(our_output - their_output	)	).item()
               __lowercase     						=torch.allclose(lowercase__,   lowercase__,   atol=1E-5	)

          print(F'''max_absolute_diff = {max_absolute_diff}'''	)  # ~ 1e-5
          print('Do both models output the same tensors?',   '🔥' if success else '💩'	)

          if not success:
               raise Exception('Something went wRoNg'	)

          if not is_folding_model:
               # Let's check contact prediction too
               __lowercase     						=model.predict_contacts(hf_tokens['input_ids'],   hf_tokens['attention_mask']	)
               __lowercase     						=esm.predict_contacts(hf_tokens['input_ids']	)
               __lowercase     						=torch.max(torch.abs(our_output - their_output	)	).item()
               __lowercase     						=torch.allclose(lowercase__,   lowercase__,   atol=1E-5	)

               print('Contact prediction testing:'	)
               print(F'''max_absolute_diff = {max_absolute_diff}'''	)  # ~ 1e-5
               print('Do both models output the same tensors?',   '🔥' if success else '💩'	)

               if not success:
                    raise Exception('Something went wRoNg'	)

          pathlib.Path(lowercase__	).mkdir(parents=lowercase__,   exist_ok=lowercase__	)
          print(F'''Saving model to {pytorch_dump_folder_path}'''	)
          model.save_pretrained(lowercase__	)

          del esm  # Free up some memory before continuing

     print(F'''Saving tokenizer to {pytorch_dump_folder_path}'''	)
     hf_tokenizer.save_pretrained(lowercase__	)

     if push_to_repo:
          model.push_to_hub(repo_id=lowercase__,   use_auth_token=lowercase__	)
          hf_tokenizer.push_to_hub(repo_id=lowercase__,   use_auth_token=lowercase__	)


if __name__ == "__main__":
       UpperCAmelCase							=					argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--pytorch_dump_folder_path", type=str, required=True, help="Path to the output PyTorch model."
       )
       parser.add_argument(
           "--classification_head", action="store_true", help="Whether to convert a final classification head."
       )
       parser.add_argument("--model", default=None, type=str, required=True, help="Name of model to convert.")
       parser.add_argument("--push_to_repo", type=str, help="Repo to upload to (including username!).")
       parser.add_argument("--auth_token", type=str, help="HuggingFace auth token.")
       UpperCAmelCase							=					parser.parse_args()
       convert_esm_checkpoint_to_pytorch(
           args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
       )





'''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


UpperCAmelCase							=					logging.get_logger(__name__)

UpperCAmelCase							=					[
    '''yjernite/retribert-base-uncased''',
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]


class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	RetriBertConfig
       lowerCAmelCase_				=	None
       lowerCAmelCase_				=	"retribert"







       def      snake_case				(							self     :	Tuple  , __lowercase     :	List[str] ):







            """simple docstring"""







            if isinstance(__lowercase  , 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(__lowercase  , 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(__lowercase  , nn.LayerNorm ):
                 module.bias.data.zero_()
                 module.weight.data.fill_(1.0 )






UpperCAmelCase							=					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 					lowerCAmelCase       ( A				):







       def __init__(							self     :	int  , __lowercase     :	RetriBertConfig ):







            """simple docstring"""







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

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

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

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







       def      snake_case				(							self     :	Dict  , __lowercase     :	List[Any]  , __lowercase     :	Any  , __lowercase     :	Optional[int]  , __lowercase     :	str=-1  , ):







            """simple docstring"""







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

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

                 # run embedding layer on everything at once
                 __lowercase     						=sent_encoder.embeddings(
                     input_ids=__lowercase  , position_ids=__lowercase  , token_type_ids=__lowercase  , inputs_embeds=__lowercase )
                 # 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(__lowercase  , __lowercase  , __lowercase )
                      pooled_output_list.append(__lowercase )
                 return torch.cat(__lowercase  , dim=0 )







       def      snake_case				(							self     :	List[Any]  , __lowercase     :	Any  , __lowercase     :	Any=None  , __lowercase     :	Optional[int]=-1  , ):







            """simple docstring"""







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







       def      snake_case				(							self     :	Dict  , __lowercase     :	Tuple  , __lowercase     :	Optional[int]=None  , __lowercase     :	Any=-1  , ):







            """simple docstring"""







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







       def      snake_case				(							self     :	int  , __lowercase     :	torch.LongTensor  , __lowercase     :	Optional[torch.FloatTensor]  , __lowercase     :	torch.LongTensor  , __lowercase     :	Optional[torch.FloatTensor]  , __lowercase     :	int = -1  , ):







            """simple docstring"""







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


