Spaces:
Sleeping
Sleeping
| """ | |
| 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, | |
| ) | |
| 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, | |
| ) | |