emrekuruu/RetrievalRouter-lambda-l50

1"""RetrievalRouter Model."""2import torch3import torch.nn as nn4from transformers import AutoConfig, PreTrainedModel, Qwen3Model5from .configuration_retrievalrouter import RetrievalRouterConfig6 7 8class RetrievalRouterModel(PreTrainedModel):9    """RAG Strategy Router - classifies queries into optimal retrieval strategies."""10    config_class = RetrievalRouterConfig11    _no_split_modules = ["Qwen3DecoderLayer"]12 13    def __init__(self, config: RetrievalRouterConfig):14        super().__init__(config)15        # Build the base architecture only; the merged base weights are loaded from this16        # checkpoint's model.safetensors by from_pretrained. Calling Qwen3Model.from_pretrained17        # here breaks under the meta-device init that from_pretrained uses.18        base_config = AutoConfig.from_pretrained(config.base_model_name)19        self.transformer = Qwen3Model(base_config)20        self.dropout = nn.Dropout(config.classifier_dropout)21        self.classifier = nn.Linear(config.hidden_size, config.num_labels)22        self.post_init()23 24    def _init_weights(self, module):25        if isinstance(module, nn.Linear):26            nn.init.normal_(module.weight, std=0.02)27            if module.bias is not None:28                nn.init.zeros_(module.bias)29 30    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):31        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)32        hidden = outputs.last_hidden_state33        if attention_mask is not None:34            mask = attention_mask.unsqueeze(-1).expand(hidden.size()).float()35            pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)36        else:37            pooled = hidden.mean(dim=1)38        logits = self.classifier(self.dropout(pooled))39        loss = self._compute_loss(logits, labels) if labels is not None else None40        return {"loss": loss, "logits": logits}41 42    def _compute_loss(self, logits, labels):43        labels_norm = labels / (labels.sum(-1, keepdim=True) + 1e-8)44        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)45        losses = -(labels_norm * log_probs).sum(-1)46        return (losses * labels.max(-1)[0]).mean()47 48    def predict(self, input_ids, attention_mask=None):49        self.eval()50        with torch.no_grad():51            logits = self.forward(input_ids, attention_mask)["logits"]52            probs = torch.softmax(logits, dim=-1)53            preds = probs.argmax(dim=-1)54        return {"predictions": preds, "probabilities": probs, 55                "strategy_names": [self.config.strategy_names[p.item()] for p in preds]}56