| """Self-contained definition of the hybrid RoBERTa + linguistic-features model. |
| |
| The released checkpoint ``hybrid_model_best.pt`` is a plain PyTorch |
| ``state_dict`` for the ``HybridClassifier`` defined here. Import this module to |
| rebuild the architecture and load the weights. |
| """ |
| import torch |
| import torch.nn as nn |
| from transformers import RobertaModel |
|
|
| CONFIG = { |
| "roberta_model": "roberta-base", |
| "roberta_dim": 768, |
| "ling_dim": 25, |
| "ling_hidden": 256, |
| "hidden_dim": 768, |
| "dropout": 0.15, |
| "freeze_bottom_layers": 6, |
| } |
|
|
| LING_FEATURE_NAMES = [ |
| "msttr", "avg_word_len", "hapax_ratio", "function_ratio", "punct_density", |
| "char_entropy", "burstiness", "repetition_ratio", "avg_sent_len", |
| "sent_len_std", "noun_ratio", "verb_ratio", "adj_ratio", "adv_ratio", |
| "pron_ratio", "pos_diversity", "avg_tree_depth", "max_tree_depth", |
| "sub_clause_ratio", "dm_density", "sent_len_cv", "fp_ratio", |
| "num_sentences", "words_per_sent", "perplexity", |
| ] |
|
|
|
|
| class FeatureAttention(nn.Module): |
| def __init__(self, num_features): |
| super().__init__() |
| self.attention = nn.Sequential( |
| nn.Linear(num_features, num_features), |
| nn.Tanh(), |
| nn.Linear(num_features, num_features), |
| nn.Softmax(dim=-1), |
| ) |
|
|
| def forward(self, features): |
| attn_weights = self.attention(features) |
| weighted_features = features * attn_weights |
| return weighted_features, attn_weights |
|
|
|
|
| class GatedFusion(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.gate = nn.Sequential( |
| nn.Linear(dim * 2, dim), |
| nn.Sigmoid(), |
| ) |
|
|
| def forward(self, cls_emb, ling_emb): |
| combined = torch.cat([cls_emb, ling_emb], dim=-1) |
| g = self.gate(combined) |
| fused = g * cls_emb + (1 - g) * ling_emb |
| return fused, g |
|
|
|
|
| class HybridClassifier(nn.Module): |
| """RoBERTa CLS embedding gated-fused with attended linguistic features.""" |
|
|
| def __init__(self, config=CONFIG): |
| super().__init__() |
| self.roberta = RobertaModel.from_pretrained(config["roberta_model"]) |
| if config["freeze_bottom_layers"] > 0: |
| for param in self.roberta.embeddings.parameters(): |
| param.requires_grad = False |
| for i in range(config["freeze_bottom_layers"]): |
| for param in self.roberta.encoder.layer[i].parameters(): |
| param.requires_grad = False |
| self.feature_attention = FeatureAttention(config["ling_dim"]) |
| self.ling_projection = nn.Sequential( |
| nn.Linear(config["ling_dim"], config["ling_hidden"]), |
| nn.LayerNorm(config["ling_hidden"]), |
| nn.GELU(), |
| nn.Dropout(config["dropout"]), |
| nn.Linear(config["ling_hidden"], config["roberta_dim"]), |
| nn.LayerNorm(config["roberta_dim"]), |
| nn.GELU(), |
| nn.Dropout(config["dropout"]), |
| ) |
| self.gated_fusion = GatedFusion(config["roberta_dim"]) |
| self.classifier = nn.Sequential( |
| nn.Linear(config["roberta_dim"], config["hidden_dim"]), |
| nn.GELU(), |
| nn.Dropout(config["dropout"]), |
| nn.Linear(config["hidden_dim"], 2), |
| ) |
| self.dropout = nn.Dropout(config["dropout"]) |
| self.config = config |
|
|
| def forward(self, input_ids, attention_mask, ling_features, return_gate=False): |
| outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask) |
| cls_embedding = outputs.last_hidden_state[:, 0, :] |
| cls_embedding = self.dropout(cls_embedding) |
| attended_features, attn_weights = self.feature_attention(ling_features) |
| ling_proj = self.ling_projection(attended_features) |
| fused, gate_values = self.gated_fusion(cls_embedding, ling_proj) |
| logits = self.classifier(fused) |
| if return_gate: |
| return logits, gate_values, attn_weights |
| return logits |
|
|
|
|
| def load_model(checkpoint_path="hybrid_model_best.pt", device="cpu"): |
| """Build the model and load the released weights. Returns an eval-mode model.""" |
| model = HybridClassifier(CONFIG) |
| state_dict = torch.load(checkpoint_path, map_location=device) |
| model.load_state_dict(state_dict) |
| model.to(device) |
| model.eval() |
| return model |
|
|
