

'''simple docstring'''

import copy
import os
from typing import Dict, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings, logging
from ..auto import AutoConfig


UpperCAmelCase							=					logging.get_logger(__name__)


UpperCAmelCase							=					{
    "suno/bark-small": "https://huggingface.co/suno/bark-small/resolve/main/config.json",
    "suno/bark": "https://huggingface.co/suno/bark/resolve/main/config.json",
}

UpperCAmelCase							=					"\n    This is the configuration class to store the configuration of a [`{model}`]. It is used to instantiate the model\n    according to the specified arguments, defining the model architecture. Instantiating a configuration with the\n    defaults will yield a similar configuration to that of the Bark [suno/bark](https://huggingface.co/suno/bark)\n    architecture.\n\n    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the\n    documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        block_size (`int`, *optional*, defaults to 1024):\n            The maximum sequence length that this model might ever be used with. Typically set this to something large\n            just in case (e.g., 512 or 1024 or 2048).\n        input_vocab_size (`int`, *optional*, defaults to 10_048):\n            Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the\n            `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with\n            regards to the chosen sub-model.\n        output_vocab_size (`int`, *optional*, defaults to 10_048):\n            Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented\n            by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought\n            with regards to the chosen sub-model.\n        num_layers (`int`, *optional*, defaults to 12):\n            Number of hidden layers in the given sub-model.\n        num_heads (`int`, *optional*, defaults to 12):\n            Number of attention heads for each attention layer in the Transformer architecture.\n        hidden_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the \"intermediate\" (often named feed-forward) layer in the architecture.\n        dropout (`float`, *optional*, defaults to 0.0):\n            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.\n        bias (`bool`, *optional*, defaults to `True`):\n            Whether or not to use bias in the linear layers and layer norm layers.\n        initializer_range (`float`, *optional*, defaults to 0.02):\n            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n"


class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	"bark_module"
       lowerCAmelCase_				=	["past_key_values"]

       lowerCAmelCase_				=	{
           "num_attention_heads": "num_heads",
           "num_hidden_layers": "num_layers",
           "vocab_size": "input_vocab_size",
           "window_size": "block_size",
       }







       def __init__(							self     :	Dict  , __lowercase     :	Dict=1024  , __lowercase     :	Any=10048  , __lowercase     :	Tuple=10048  , __lowercase     :	Optional[Any]=12  , __lowercase     :	List[Any]=12  , __lowercase     :	Any=768  , __lowercase     :	Union[str, Any]=0.0  , __lowercase     :	Union[str, Any]=True  , __lowercase     :	Optional[int]=0.0_2  , __lowercase     :	Union[str, Any]=True  , **__lowercase     :	Optional[Any]  , ):







            """simple docstring"""







            __lowercase     						=block_size
            __lowercase     						=input_vocab_size
            __lowercase     						=output_vocab_size
            __lowercase     						=num_layers
            __lowercase     						=num_heads
            __lowercase     						=hidden_size
            __lowercase     						=dropout
            __lowercase     						=bias
            __lowercase     						=use_cache
            __lowercase     						=initializer_range

            super().__init__(**__lowercase )







       @classmethod
       def      snake_case				(							cls     :	str  , __lowercase     :	Union[str, os.PathLike]  , __lowercase     :	Optional[Union[str, os.PathLike]] = None  , __lowercase     :	bool = False  , __lowercase     :	bool = False  , __lowercase     :	Optional[Union[str, bool]] = None  , __lowercase     :	str = "main"  , **__lowercase     :	List[Any]  , ):







            """simple docstring"""







            __lowercase     						=cache_dir
            __lowercase     						=force_download
            __lowercase     						=local_files_only
            __lowercase     						=revision

            cls._set_token_in_kwargs(__lowercase  , __lowercase )

            __lowercase						,       __lowercase     						=cls.get_config_dict(__lowercase  , **__lowercase )

            # get the config dict if we are loading from Bark
            if config_dict.get('model_type' ) == "bark":
                 __lowercase     						=config_dict[f'''{cls.model_type}_config''']

            if "model_type" in config_dict and hasattr(cls  , 'model_type' ) and config_dict["model_type"] != cls.model_type:
                 logger.warning(
                     f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
                     f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.''' )

            return cls.from_dict(__lowercase  , **__lowercase )






@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config="BarkSemanticConfig"				,     model="BarkSemanticModel"				)				,     "\n    Example:\n\n    ```python\n    >>> from transformers import BarkSemanticConfig, BarkSemanticModel\n\n    >>> # Initializing a Bark sub-module style configuration\n    >>> configuration = BarkSemanticConfig()\n\n    >>> # Initializing a model (with random weights) from the suno/bark style configuration\n    >>> model = BarkSemanticModel(configuration)\n\n    >>> # Accessing the model configuration\n    >>> configuration = model.config\n    ```"				,     )
class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	"semantic"






@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config="BarkCoarseConfig"				,     model="BarkCoarseModel"				)				,     "\n    Example:\n\n    ```python\n    >>> from transformers import BarkCoarseConfig, BarkCoarseModel\n\n    >>> # Initializing a Bark sub-module style configuration\n    >>> configuration = BarkCoarseConfig()\n\n    >>> # Initializing a model (with random weights) from the suno/bark style configuration\n    >>> model = BarkCoarseModel(configuration)\n\n    >>> # Accessing the model configuration\n    >>> configuration = model.config\n    ```"				,     )
class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	"coarse_acoustics"






@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config="BarkFineConfig"				,     model="BarkFineModel"				)				,     "\n        n_codes_total (`int`, *optional*, defaults to 8):\n            The total number of audio codebooks predicted. Used in the fine acoustics sub-model.\n        n_codes_given (`int`, *optional*, defaults to 1):\n            The number of audio codebooks predicted in the coarse acoustics sub-model. Used in the acoustics\n            sub-models.\n    Example:\n\n    ```python\n    >>> from transformers import BarkFineConfig, BarkFineModel\n\n    >>> # Initializing a Bark sub-module style configuration\n    >>> configuration = BarkFineConfig()\n\n    >>> # Initializing a model (with random weights) from the suno/bark style configuration\n    >>> model = BarkFineModel(configuration)\n\n    >>> # Accessing the model configuration\n    >>> configuration = model.config\n    ```"				,     )
class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	"fine_acoustics"







       def __init__(							self     :	Union[str, Any]  , __lowercase     :	Any=True  , __lowercase     :	Optional[Any]=8  , __lowercase     :	str=1  , **__lowercase     :	str ):







            """simple docstring"""







            __lowercase     						=n_codes_total
            __lowercase     						=n_codes_given

            super().__init__(tie_word_embeddings=__lowercase  , **__lowercase )






class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	"bark"
       lowerCAmelCase_				=	True







       def __init__(							self     :	str  , __lowercase     :	Dict = None  , __lowercase     :	Dict = None  , __lowercase     :	Dict = None  , __lowercase     :	Dict = None  , __lowercase     :	List[Any]=0.0_2  , **__lowercase     :	Tuple  , ):







            """simple docstring"""







            if semantic_config is None:
                 __lowercase     						={}
                 logger.info('semantic_config is None. initializing the semantic model with default values.' )

            if coarse_acoustics_config is None:
                 __lowercase     						={}
                 logger.info('coarse_acoustics_config is None. initializing the coarse model with default values.' )

            if fine_acoustics_config is None:
                 __lowercase     						={}
                 logger.info('fine_acoustics_config is None. initializing the fine model with default values.' )

            if codec_config is None:
                 __lowercase     						={}
                 logger.info('codec_config is None. initializing the codec model with default values.' )

            __lowercase     						=BarkSemanticConfig(**__lowercase )
            __lowercase     						=BarkCoarseConfig(**__lowercase )
            __lowercase     						=BarkFineConfig(**__lowercase )
            __lowercase     						=AutoConfig.for_model(**__lowercase )

            __lowercase     						=initializer_range

            super().__init__(**__lowercase )







       @classmethod
       def      snake_case				(							cls     :	List[str]  , __lowercase     :	BarkSemanticConfig  , __lowercase     :	BarkCoarseConfig  , __lowercase     :	BarkFineConfig  , __lowercase     :	AutoConfig  , **__lowercase     :	List[Any]  , ):







            """simple docstring"""







            return cls(
                semantic_config=semantic_config.to_dict()  , coarse_acoustics_config=coarse_acoustics_config.to_dict()  , fine_acoustics_config=fine_acoustics_config.to_dict()  , codec_config=codec_config.to_dict()  , **__lowercase  , )







       def      snake_case				(							self     :	Any ):







            """simple docstring"""







            __lowercase     						=copy.deepcopy(self.__dict__ )

            __lowercase     						=self.semantic_config.to_dict()
            __lowercase     						=self.coarse_acoustics_config.to_dict()
            __lowercase     						=self.fine_acoustics_config.to_dict()
            __lowercase     						=self.codec_config.to_dict()

            __lowercase     						=self.__class__.model_type
            return output





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


