"""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())