File size: 1,126 Bytes
1b60a2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | from typing import Any
from transformers import AutoConfig, PretrainedConfig
class EmbedderConfig(PretrainedConfig):
model_type = "embedder"
def __init__(
self,
base_model_name: str = "nomic-ai/nomic-embed-text-v1.5",
num_blocks: int = 2,
dropout: float = 0.0,
encoder_only: bool = False,
**kwargs: Any,
):
super().__init__(**kwargs)
self.base_model_name = base_model_name
self.num_blocks = num_blocks
self.dropout = dropout
self.encoder_only = encoder_only
self.encoder_config = AutoConfig.from_pretrained( # nosec B615
base_model_name,
trust_remote_code=True,
)
self.auto_map = {
"AutoConfig": "configuration_embedder.EmbedderConfig",
"AutoModel": "modeling_embedder.EmbedderModel",
}
def __getattr__(self, name: str) -> Any:
if name != "encoder_config" and hasattr(self.encoder_config, name):
return getattr(self.encoder_config, name)
raise AttributeError(name)
EmbedderConfig.register_for_auto_class()
|