test / configuration_embedder.py
JalalKhal's picture
hf model and vllm integration files
c3443ec verified
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(
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()