emrekuruu/RetrievalRouter-lambda-l50
066
1"""RetrievalRouter Configuration."""2from transformers import PretrainedConfig3 4# Standalone copy of train.config.ARM_NAMES: this module is uploaded to the Hub and loaded5# via trust_remote_code, so it cannot import from the training package.6STRATEGY_NAMES = ["MULTIMODAL_RERANK", "MULTIMODAL-SINGLE", "TEXT_RERANK", "TEXT-SINGLE", "BM25"]7 8 9class RetrievalRouterConfig(PretrainedConfig):10 """Configuration for RetrievalRouter - a query-aware retrieval router."""11 model_type = "retrievalrouter"12 13 def __init__(14 self,15 base_model_name: str = "Qwen/Qwen3-0.6B-Base",16 hidden_size: int = 1024,17 num_labels: int = 5,18 classifier_dropout: float = 0.1,19 strategy_names: list = None,20 **kwargs,21 ):22 super().__init__(num_labels=num_labels, **kwargs)23 self.base_model_name = base_model_name24 self.hidden_size = hidden_size25 self.classifier_dropout = classifier_dropout26 self.strategy_names = strategy_names or STRATEGY_NAMES27 