File size: 7,351 Bytes
f829eea 7bafcf0 f829eea 3f30f83 f829eea 3f30f83 f829eea 3f30f83 f829eea 3f30f83 f829eea | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """BetterWright task-conditioned relevance heads on LFM2.5 Encoder."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.modeling_outputs import ModelOutput
from transformers.models.lfm2.configuration_lfm2 import Lfm2Config
from transformers.models.lfm2.modeling_lfm2 import Lfm2PreTrainedModel
from .modeling_lfm2_bidirectional import Lfm2BidirectionalModel, _install_patches
@dataclass
class BetterWrightEncoderOutput(ModelOutput):
loss: Optional[torch.Tensor] = None
logits: Optional[torch.Tensor] = None
token_logits: Optional[torch.Tensor] = None
uncertainty_logits: Optional[torch.Tensor] = None
token_loss: Optional[torch.Tensor] = None
token_rank_loss: Optional[torch.Tensor] = None
last_hidden_state: Optional[torch.Tensor] = None
class BetterWrightEncoder(Lfm2PreTrainedModel):
"""Cross-encoder chunk scorer plus per-token relevance and fallback heads."""
config_class = Lfm2Config
base_model_prefix = "lfm2"
def __init__(self, config: Lfm2Config):
_install_patches()
config.use_cache = False
super().__init__(config)
self.lfm2 = Lfm2BidirectionalModel(config)
self.relevance_head = nn.Sequential(
nn.Linear(config.hidden_size, config.hidden_size // 2),
nn.SiLU(),
nn.Dropout(0.05),
nn.Linear(config.hidden_size // 2, 1),
)
self.token_head = nn.Linear(config.hidden_size, 1)
self.uncertainty_head = nn.Linear(config.hidden_size, 1)
self.post_init()
def get_input_embeddings(self):
return self.lfm2.embed_tokens
def set_input_embeddings(self, value):
self.lfm2.embed_tokens = value
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.Tensor] = None,
token_labels: Optional[torch.Tensor] = None,
return_dict: bool = True,
**kwargs,
) -> BetterWrightEncoderOutput | tuple:
outputs = self.lfm2(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
inputs_embeds=inputs_embeds,
use_cache=False,
return_dict=True,
**kwargs,
)
hidden = outputs.last_hidden_state
token_logits = self.token_head(hidden).squeeze(-1)
if attention_mask is None:
pooling_weights = torch.softmax(token_logits, dim=1).unsqueeze(-1)
else:
masked_pool_logits = token_logits.masked_fill(attention_mask == 0, -1e4)
pooling_weights = torch.softmax(masked_pool_logits, dim=1).unsqueeze(-1)
# Learned token pooling lets pair-level relevance supervision teach the
# token head where evidence lives before the exact-span curriculum.
pooled = (hidden * pooling_weights.to(hidden.dtype)).sum(dim=1)
logits = self.relevance_head(pooled).squeeze(-1)
uncertainty_logits = self.uncertainty_head(pooled).squeeze(-1)
loss = None
token_loss = None
token_rank_loss = None
if labels is not None:
labels = labels.to(logits.dtype)
positive_weight = float(getattr(self.config, "betterwright_positive_weight", 3.0))
relevance_loss = F.binary_cross_entropy_with_logits(
logits,
labels,
pos_weight=torch.tensor(positive_weight, device=logits.device, dtype=logits.dtype),
)
# Predict disagreement/ambiguity most strongly around the decision boundary.
uncertainty_target = 1.0 - (labels - torch.sigmoid(logits).detach()).abs()
uncertainty_loss = F.binary_cross_entropy_with_logits(
uncertainty_logits, uncertainty_target.clamp(0, 1)
)
loss = relevance_loss + 0.08 * uncertainty_loss
if token_labels is not None:
valid = token_labels >= 0
if valid.any():
valid_labels = token_labels[valid]
positive_count = (valid_labels == 1).sum()
negative_count = (valid_labels == 0).sum()
if positive_count:
# Required lines are exceptionally sparse in a 64K tree.
# Batch-adaptive weighting prevents the all-negative token
# solution while remaining capped for stable BF16 updates.
token_positive_weight = float(
(negative_count / positive_count).clamp(20, 512)
)
else:
token_positive_weight = 1.0
token_loss = F.binary_cross_entropy_with_logits(
token_logits[valid],
valid_labels.to(token_logits.dtype),
pos_weight=torch.tensor(
token_positive_weight,
device=token_logits.device,
dtype=token_logits.dtype,
),
)
# Runtime pruning ranks lines by their highest-scoring token.
# BCE alone can achieve a reasonable average while leaving a
# few distractor tokens above the evidence. Explicitly compare
# each positive row against its hardest negative tokens so the
# training objective matches that production ranking behavior.
token_rank_losses = []
for row_index in range(token_logits.shape[0]):
row_valid = valid[row_index]
row_labels = token_labels[row_index]
positives = token_logits[row_index][row_valid & (row_labels == 1)]
negatives = token_logits[row_index][row_valid & (row_labels == 0)]
if positives.numel() == 0 or negatives.numel() == 0:
continue
hard_count = min(64, negatives.numel())
hard_negatives = torch.topk(
negatives.float(), hard_count, sorted=False
).values
token_rank_losses.append(
F.softplus(
1.0 + hard_negatives.mean() - positives.float().mean()
)
)
if token_rank_losses:
token_rank_loss = torch.stack(token_rank_losses).mean()
else:
token_rank_loss = token_loss.new_zeros(())
token_objective = 0.75 * token_loss + token_rank_loss
loss = token_objective if loss is None else loss + token_objective
result = BetterWrightEncoderOutput(
loss=loss,
logits=logits,
token_logits=token_logits,
uncertainty_logits=uncertainty_logits,
token_loss=token_loss,
token_rank_loss=token_rank_loss,
last_hidden_state=hidden,
)
return result if return_dict else tuple(result.values())
|