codingutf8 2020 the rag s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license rag model configuration from configurationutils import pretrainedconfig from utils import addstartdocstrings ragconfigdoc r ragconfig stores the configuration of a ragmodel configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information args titlesep str optional defaults to separator inserted between the title and the text of the retrieved document when calling ragretriever docsep str optional defaults to separator inserted between the text of the retrieved document and the original input when calling ragretriever ndocs int optional defaults to 5 number of documents to retrieve maxcombinedlength int optional defaults to 300 max length of contextualized input returned by ragretriever call retrievalvectorsize int optional defaults to 768 dimensionality of the document embeddings indexed by ragretriever retrievalbatchsize int optional defaults to 8 retrieval batch size defined as the number of queries issues concurrently to the faiss index encapsulated ragretriever dataset str optional defaults to wikidpr a dataset identifier of the indexed dataset in huggingface datasets list all available datasets and ids using datasets listdatasets datasetsplit str optional defaults to train which split of the dataset to load indexname str optional defaults to compressed the index name of the index associated with the dataset one can choose between legacy exact and compressed indexpath str optional the path to the serialized faiss index on disk passagespath str optional a path to text passages compatible with the faiss index required if using models rag retrievalrag legacyindex usedummydataset bool optional defaults to false whether to load a dummy variant of the dataset specified by dataset labelsmoothing float optional defaults to 0 0 only relevant if returnloss is set to true controls the epsilon parameter value for label smoothing in the loss calculation if set to 0 no label smoothing is performed domarginalize bool optional defaults to false if true the logits are marginalized over all documents by making use of torch nn functional logsoftmax reduceloss bool optional defaults to false whether or not to reduce the nll loss using the torch tensor sum operation dodeduplication bool optional defaults to true whether or not to deduplicate the generations from different context documents for a given input has to be set to false if used while training with distributed backend excludebosscore bool optional defaults to false whether or not to disregard the bos token when computing the loss outputretrievedbool optional defaults to false if set to true retrieveddocembeds retrieveddocids contextinputids and contextattentionmask are returned see returned tensors for more detail usecache bool optional defaults to true whether or not the model should return the last keyvalues attentions not used by all models forcedeostokenid int optional the id of the token to force as the last generated token when maxlength is reached usually set to eostokenid return clsquestionencoderquestionencoderconfig todict generatorgeneratorconfig todict kwargs coding utf 8 2020 the rag s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license rag model configuration ragconfig stores the configuration of a ragmodel configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information args title_sep str optional defaults to separator inserted between the title and the text of the retrieved document when calling ragretriever doc_sep str optional defaults to separator inserted between the text of the retrieved document and the original input when calling ragretriever n_docs int optional defaults to 5 number of documents to retrieve max_combined_length int optional defaults to 300 max length of contextualized input returned by ragretriever __call__ retrieval_vector_size int optional defaults to 768 dimensionality of the document embeddings indexed by ragretriever retrieval_batch_size int optional defaults to 8 retrieval batch size defined as the number of queries issues concurrently to the faiss index encapsulated ragretriever dataset str optional defaults to wiki_dpr a dataset identifier of the indexed dataset in huggingface datasets list all available datasets and ids using datasets list_datasets dataset_split str optional defaults to train which split of the dataset to load index_name str optional defaults to compressed the index name of the index associated with the dataset one can choose between legacy exact and compressed index_path str optional the path to the serialized faiss index on disk passages_path str optional a path to text passages compatible with the faiss index required if using models rag retrieval_rag legacyindex use_dummy_dataset bool optional defaults to false whether to load a dummy variant of the dataset specified by dataset label_smoothing float optional defaults to 0 0 only relevant if return_loss is set to true controls the epsilon parameter value for label smoothing in the loss calculation if set to 0 no label smoothing is performed do_marginalize bool optional defaults to false if true the logits are marginalized over all documents by making use of torch nn functional log_softmax reduce_loss bool optional defaults to false whether or not to reduce the nll loss using the torch tensor sum operation do_deduplication bool optional defaults to true whether or not to deduplicate the generations from different context documents for a given input has to be set to false if used while training with distributed backend exclude_bos_score bool optional defaults to false whether or not to disregard the bos token when computing the loss output_retrieved bool optional defaults to false if set to true retrieved_doc_embeds retrieved_doc_ids context_input_ids and context_attention_mask are returned see returned tensors for more detail use_cache bool optional defaults to true whether or not the model should return the last key values attentions not used by all models forced_eos_token_id int optional the id of the token to force as the last generated token when max_length is reached usually set to eos_token_id instantiate a encoderdecoderconfig or a derived class from a pre trained encoder model configuration and decoder model configuration returns encoderdecoderconfig an instance of a configuration object
from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


RAG_CONFIG_DOC = r


@add_start_docstrings(RAG_CONFIG_DOC)
class RagConfig(PretrainedConfig):
    model_type = "rag"
    is_composition = True

    def __init__(
        self,
        vocab_size=None,
        is_encoder_decoder=True,
        prefix=None,
        bos_token_id=None,
        pad_token_id=None,
        eos_token_id=None,
        decoder_start_token_id=None,
        title_sep=" / ",
        doc_sep=" // ",
        n_docs=5,
        max_combined_length=300,
        retrieval_vector_size=768,
        retrieval_batch_size=8,
        dataset="wiki_dpr",
        dataset_split="train",
        index_name="compressed",
        index_path=None,
        passages_path=None,
        use_dummy_dataset=False,
        reduce_loss=False,
        label_smoothing=0.0,
        do_deduplication=True,
        exclude_bos_score=False,
        do_marginalize=False,
        output_retrieved=False,
        use_cache=True,
        forced_eos_token_id=None,
        **kwargs,
    ):
        super().__init__(
            bos_token_id=bos_token_id,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            decoder_start_token_id=decoder_start_token_id,
            forced_eos_token_id=forced_eos_token_id,
            is_encoder_decoder=is_encoder_decoder,
            prefix=prefix,
            vocab_size=vocab_size,
            **kwargs,
        )
        assert (
            "question_encoder" in kwargs and "generator" in kwargs
        ), "Config has to be initialized with question_encoder and generator config"
        question_encoder_config = kwargs.pop("question_encoder")
        question_encoder_model_type = question_encoder_config.pop("model_type")
        decoder_config = kwargs.pop("generator")
        decoder_model_type = decoder_config.pop("model_type")

        from ..auto.configuration_auto import AutoConfig

        self.question_encoder = AutoConfig.for_model(question_encoder_model_type, **question_encoder_config)
        self.generator = AutoConfig.for_model(decoder_model_type, **decoder_config)

        self.reduce_loss = reduce_loss
        self.label_smoothing = label_smoothing
        self.exclude_bos_score = exclude_bos_score
        self.do_marginalize = do_marginalize

        self.title_sep = title_sep
        self.doc_sep = doc_sep
        self.n_docs = n_docs
        self.max_combined_length = max_combined_length

        self.dataset = dataset
        self.dataset_split = dataset_split
        self.index_name = index_name

        self.retrieval_vector_size = retrieval_vector_size
        self.retrieval_batch_size = retrieval_batch_size
        self.passages_path = passages_path
        self.index_path = index_path
        self.use_dummy_dataset = use_dummy_dataset

        self.output_retrieved = output_retrieved

        self.do_deduplication = do_deduplication

        self.use_cache = use_cache

        if self.forced_eos_token_id is None:
            self.forced_eos_token_id = getattr(self.generator, "forced_eos_token_id", None)

    @classmethod
    def from_question_encoder_generator_configs(
        cls, question_encoder_config: PretrainedConfig, generator_config: PretrainedConfig, **kwargs
    ) -> PretrainedConfig:
        r
        return cls(question_encoder=question_encoder_config.to_dict(), generator=generator_config.to_dict(), **kwargs)