| import math |
| from typing import Any, Dict, Optional |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import Qwen3ForCausalLM |
|
|
|
|
| class QueritModel(Qwen3ForCausalLM): |
| """Querit reranker based on Qwen3-Embedding-4B backbone with binary classification head.""" |
|
|
| def __init__(self, config, use_lm_head: bool = False): |
| super().__init__(config) |
| hidden_size = self.config.hidden_size |
| self.head = nn.Linear(hidden_size, 2) |
| nn.init.normal_(self.head.weight, std=1e-4) |
| nn.init.zeros_(self.head.bias) |
| if not use_lm_head: |
| self.lm_head = None |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| labels: Optional[torch.Tensor] = None, |
| scores: Optional[torch.Tensor] = None, |
| qids: Optional[torch.Tensor] = None, |
| ) -> Dict[str, Any]: |
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| ) |
| cls_hidden = outputs.last_hidden_state[:, -1, :] |
| logits = self.head(cls_hidden) |
| probs = torch.softmax(logits, dim=-1) |
| pred_labels = torch.argmax(probs, dim=-1) |
| rank_scores = self._compute_score(probs) |
|
|
| loss = None |
| if labels is not None and scores is not None: |
| loss = self._pairwise_hinge_loss(rank_scores, scores, qids) |
|
|
| return { |
| "loss": loss, |
| "qids": qids, |
| "score": rank_scores, |
| "pred_label": pred_labels, |
| } |
|
|
| def _pairwise_hinge_loss( |
| self, |
| logits: torch.Tensor, |
| labels: torch.Tensor, |
| qids: torch.Tensor, |
| margin_weight: float = 0.8, |
| topk: bool = False, |
| pairdiff_mask: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| qid_mask = (qids.unsqueeze(0) == qids.unsqueeze(1)).float() |
| if topk: |
| qid_mask = qid_mask * self._get_topk_mask(qids, logits.squeeze(-1), labels) |
| batch_size = logits.shape[0] |
| labels = labels.unsqueeze(1) |
|
|
| score_pos = logits.expand(-1, batch_size) |
| score_neg = score_pos.transpose(0, 1) |
| pos = labels.expand(-1, batch_size) |
| neg = pos.transpose(0, 1) |
|
|
| if pairdiff_mask is not None: |
| margin = (pos - neg + pairdiff_mask) * qid_mask * margin_weight |
| else: |
| margin = (pos - neg) * qid_mask * margin_weight |
| pair_mask = (margin > 1e-6).float() |
|
|
| score_diff = score_pos - score_neg |
| margin_diff = margin + torch.clamp(-score_diff, min=-10.0) |
| loss = torch.relu(margin_diff) * pair_mask |
| return torch.sum(loss) / (torch.sum(pair_mask) + 1e-5) |
|
|
| def _get_topk_mask(self, qids, logits, labels): |
| flatten_qids = qids.view(-1) |
| flatten_logits = logits.view(-1) |
| flatten_labels = labels.view(-1) |
| unique_qids = torch.unique(flatten_qids) |
|
|
| batch_size = qids.shape[0] |
| position_mask = torch.ones(batch_size, dtype=torch.float32, device=logits.device) |
| for uq in unique_qids: |
| indices = (flatten_qids == uq).nonzero(as_tuple=True)[0] |
| cur_labels = flatten_labels[indices] |
| valid_count = (cur_labels >= 0).sum().item() |
| k = math.ceil(valid_count * 0.3) |
| if k == 0: |
| continue |
| cur_logits = flatten_logits[indices] |
| topk_idx = indices[cur_logits.argsort(descending=True)[:k]] |
| position_mask[topk_idx] = 2.0 |
|
|
| return position_mask.unsqueeze(-1).expand(batch_size, batch_size).transpose(0, 1) |
|
|
| def _compute_score(self, probs: torch.Tensor) -> torch.Tensor: |
| weights = torch.tensor([-1.0, 1.0], device=probs.device) |
| return (probs * weights).sum(dim=-1, keepdim=True) |
|
|
