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