File size: 3,770 Bytes
5fade51 | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | 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)
|