| """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) | |