"""HuggingFace config for DiffRetriever `trust_remote_code` loading.

A thin `PretrainedConfig` so that

    AutoModel.from_pretrained("ielabgroup/diffretriever-...", trust_remote_code=True)

can route to `DiffRetrieverModel` via the repo's `config.json` `auto_map`.

The real retrieval configuration (prompt token ids, K_q/K_p, temperature,
sparse weight, ...) lives in `retriever_config.json` and is read by
`TrainableDiffusionRetriever.load()`. The fields here are informational only
(they show up in the Hub config viewer) and are not required for loading.

This file is shipped *inside each model repo* — keep it dependency-light.
"""
from transformers import PretrainedConfig


class DiffRetrieverConfig(PretrainedConfig):
    model_type = "diffretriever"

    def __init__(
        self,
        base_model: str | None = None,
        backbone_type: str | None = None,
        mode: str = "single",
        k_q: int = 1,
        k_p: int = 1,
        **kwargs,
    ):
        self.base_model = base_model      # e.g. "Dream-org/Dream-v0-Instruct-7B"
        self.backbone_type = backbone_type  # e.g. "dream" / "llada"
        self.mode = mode                  # "single" | "multi"
        self.k_q = k_q
        self.k_p = k_p
        super().__init__(**kwargs)