nvidia

/

NV-Embed-v1

like 397

Hi there.
This may not be necessary but I think we should provide some features, for example, hidden size in config class to apply for various libraries.

class NVEmbedConfig(PretrainedConfig):
    model_type = "nvembed"
    is_composition = False

    def __init__(
        self,
        latent_attention_config=None,
        text_config=None,
        padding_side: Literal["right", "left"]="right",
        add_pad_token: bool=True,
        is_mask_instruction: bool = True,
        add_eos: bool=True,
        mask_type: str="b",
        **kwargs,
    ):
        if isinstance(latent_attention_config, dict):
            latent_attention_config["model_type"] = (
                latent_attention_config["model_type"] if "model_type" in latent_attention_config else LATENT_ATTENTION_TYPE
            )
            latent_attention_config = CONFIG_MAPPING[latent_attention_config["model_type"]](**latent_attention_config)
        elif latent_attention_config is None:
            latent_attention_config = CONFIG_MAPPING[LATENT_ATTENTION_TYPE]()

        self.latent_attention_config = latent_attention_config

        if isinstance(text_config, dict):
            text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama"
            text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
        elif text_config is None:
            text_config = None

        self.text_config = text_config
        self.padding_side = padding_side
        self.is_mask_instruction = is_mask_instruction
        self.add_pad_token = add_pad_token
        self.add_eos = add_eos
        self.mask_type = mask_type
        if "hidden_size" in kwargs:
            self.hidden_size = kwargs["hidden_size"]
        else:
            self.hidden_size = 4096

        super().__init__(**kwargs)