configuration_chest2vec.py

"""Configuration for Chest2Vec — a Qwen3-Embedding model (LoRA merged in) for chest
radiology report embeddings.

Self-contained: the merged encoder weights ship in the repo, so loading needs neither the
`chest2vec` package nor a download of the base Qwen3-Embedding weights. Produces a single
L2-normalized report embedding (last-token / EOS pooling), Qwen3-Embedding convention.
"""
from typing import Optional
from transformers import PretrainedConfig


class Chest2VecConfig(PretrainedConfig):
    model_type = "chest2vec"

    def __init__(
        self,
        encoder_config: Optional[dict] = None,
        base_model: str = "Qwen/Qwen3-Embedding-0.6B",
        hidden_size: int = 1024,
        default_max_len: int = 512,
        pooling: str = "last_token",
        matryoshka_dims: Optional[list] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.encoder_config = encoder_config or {}
        self.base_model = base_model
        self.hidden_size = hidden_size
        self.default_max_len = default_max_len
        self.pooling = pooling
        self.matryoshka_dims = matryoshka_dims or []