Spaces:
Runtime error
Runtime error
File size: 8,089 Bytes
fc7db98 | 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 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | """
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,
)
|