""" Stage E — Two-head model for 5-class word complexity + 3-class reason. No regression head. Outputs are discrete levels (Very Easy → Very Hard) and, when Hard/Very Hard, one of three explainable difficulty causes. Novel architecture (publish contribution): one encoder, dual heads, masked reason loss, target-span pooling for word-in-context LCP. """ 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 import AutoModel, AutoTokenizer, PreTrainedTokenizer from linguistic_features import N_LINGUISTIC_FEATURES from utils import ( MODELS, POOLING_CLS_CONCAT, POOLING_CLS_ONLY, POOLING_SPAN, POOLING_TGT_MARKER, TGT_END_TOKEN, TGT_TOKEN, ) @dataclass class ModelOutput: level_logits: torch.Tensor reason_logits: torch.Tensor pooled: torch.Tensor class TwoHeadModel(nn.Module): """Shared encoder with level (5-class) and reason (3-class) heads.""" def __init__( self, model_name: str = MODELS["deberta"], hidden_dropout: float = 0.1, pooling_mode: str = POOLING_SPAN, use_linguistic_features: bool = False, # Legacy alias: cls_concat maps to cls_concat pooling use_cls_concat: bool | None = None, ): super().__init__() if use_cls_concat is not None: pooling_mode = POOLING_CLS_CONCAT if use_cls_concat else POOLING_SPAN self.encoder = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32) hidden = self.encoder.config.hidden_size self.pooling_mode = pooling_mode self.use_linguistic_features = use_linguistic_features self.model_name = model_name extra = N_LINGUISTIC_FEATURES if use_linguistic_features else 0 if pooling_mode == POOLING_CLS_CONCAT: head_in = hidden * 2 + extra else: head_in = hidden + extra self.dropout = nn.Dropout(hidden_dropout) self.level_head = nn.Linear(head_in, 5) self.reason_head = nn.Linear(head_in, 3) self._tgt_id: Optional[int] = None self._tgt_end_id: Optional[int] = None def set_tgt_token_ids(self, tgt_id: int, tgt_end_id: int | None = None) -> None: self._tgt_id = tgt_id self._tgt_end_id = tgt_end_id def set_tgt_token_id(self, tgt_id: int) -> None: self.set_tgt_token_ids(tgt_id, self._tgt_end_id) def _span_pool(self, last_hidden: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: """Mean-pool token hidden states inside [TGT] ... [/TGT] (target word subwords).""" batch_size = input_ids.size(0) pooled = [] cls_vec = last_hidden[:, 0, :] for i in range(batch_size): ids = input_ids[i] open_pos = (ids == self._tgt_id).nonzero(as_tuple=True)[0] if self._tgt_id is not None else None if open_pos is None or len(open_pos) == 0: pooled.append(cls_vec[i]) continue start = int(open_pos[0].item()) + 1 end = len(ids) if self._tgt_end_id is not None: close_pos = (ids == self._tgt_end_id).nonzero(as_tuple=True)[0] close_after = close_pos[close_pos > open_pos[0]] if len(close_after) > 0: end = int(close_after[0].item()) span_idx = [j for j in range(start, end) if ids[j].item() != 0] # skip pad if not span_idx: pooled.append(last_hidden[i, int(open_pos[0].item()), :]) else: vecs = last_hidden[i, span_idx, :] pooled.append(vecs.mean(dim=0)) return torch.stack(pooled, dim=0) def _first_marker_pool(self, last_hidden: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: cls_vec = last_hidden[:, 0, :] if self._tgt_id is None: return cls_vec batch_size = input_ids.size(0) vecs = [] for i in range(batch_size): positions = (input_ids[i] == self._tgt_id).nonzero(as_tuple=True)[0] if len(positions) > 0: vecs.append(last_hidden[i, positions[0], :]) else: vecs.append(cls_vec[i]) return torch.stack(vecs, dim=0) def _pool(self, last_hidden: torch.Tensor, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: cls_vec = last_hidden[:, 0, :] if self.pooling_mode == POOLING_CLS_ONLY: return cls_vec if self.pooling_mode == POOLING_SPAN: return self._span_pool(last_hidden, input_ids) if self.pooling_mode == POOLING_TGT_MARKER: return self._first_marker_pool(last_hidden, input_ids) if self.pooling_mode == POOLING_CLS_CONCAT: tgt_vec = self._first_marker_pool(last_hidden, input_ids) return torch.cat([cls_vec, tgt_vec], dim=-1) return self._span_pool(last_hidden, input_ids) def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor, linguistic_features: Optional[torch.Tensor] = None, ) -> ModelOutput: outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) pooled = self._pool(outputs.last_hidden_state, input_ids, attention_mask) if self.use_linguistic_features: if linguistic_features is None: raise ValueError("linguistic_features required when use_linguistic_features=True") pooled = torch.cat([pooled, linguistic_features], dim=-1) pooled = self.dropout(pooled) # Encoder may run in float16 on GPU; classification heads stay float32. pooled = pooled.to(dtype=self.level_head.weight.dtype) return ModelOutput( level_logits=self.level_head(pooled), reason_logits=self.reason_head(pooled), pooled=pooled, ) def add_tgt_tokens(tokenizer: PreTrainedTokenizer) -> tuple[int, int]: """Add [TGT] and [/TGT] special tokens; return their ids.""" special = {"additional_special_tokens": [TGT_TOKEN, TGT_END_TOKEN]} tokenizer.add_special_tokens(special) return ( tokenizer.convert_tokens_to_ids(TGT_TOKEN), tokenizer.convert_tokens_to_ids(TGT_END_TOKEN), ) def add_tgt_token(tokenizer: PreTrainedTokenizer) -> int: """Backward-compatible: add markers and return open [TGT] id.""" open_id, _ = add_tgt_tokens(tokenizer) return open_id def compute_loss( level_logits: torch.Tensor, reason_logits: torch.Tensor, level_ids: torch.Tensor, reason_ids: torch.Tensor, reason_mask: torch.Tensor, reason_class_weights: Optional[torch.Tensor] = None, lambda_reason: float = 1.0, level_only: bool = False, ) -> tuple[torch.Tensor, dict]: level_loss = F.cross_entropy(level_logits, level_ids) if level_only: return level_loss, {"level_loss": level_loss.item(), "reason_loss": 0.0, "total_loss": level_loss.item()} per_row = F.cross_entropy( reason_logits, reason_ids, weight=reason_class_weights, reduction="none", ) masked = per_row * reason_mask denom = reason_mask.sum().clamp(min=1.0) reason_loss = masked.sum() / denom total = level_loss + lambda_reason * reason_loss return total, { "level_loss": level_loss.item(), "reason_loss": reason_loss.item(), "total_loss": total.item(), } def load_tokenizer(model_key: str = "deberta") -> PreTrainedTokenizer: return AutoTokenizer.from_pretrained(MODELS[model_key]) def build_model( model_key: str = "deberta", pooling_mode: str = POOLING_SPAN, use_linguistic_features: bool = False, use_cls_concat: bool | None = None, ) -> TwoHeadModel: return TwoHeadModel( model_name=MODELS[model_key], pooling_mode=pooling_mode, use_linguistic_features=use_linguistic_features, use_cls_concat=use_cls_concat, )