File size: 2,657 Bytes
431bcf6 | 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 | # ─────────────────────────────────────────────────────────────────────────────
# Fine-Tuned BERT Regressor Architecture
# src/models/bert_regressor.py
# ─────────────────────────────────────────────────────────────────────────────
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
class BertRegressorModel(nn.Module):
"""
End-to-end fine-tuned BERT Regression model for Extraversion prediction.
Architecture:
bert-base-uncased → [CLS] hidden state (768-dim)
→ Dropout(0.3)
→ Linear(768, 256)
→ ReLU
→ Dropout(0.2)
→ Linear(256, 1)
→ Sigmoid * 99 (output in 0–99 range)
"""
def __init__(self, dropout1: float = 0.3, dropout2: float = 0.2):
super().__init__()
self.bert = BertModel.from_pretrained("bert-base-uncased")
hidden = self.bert.config.hidden_size # 768
self.regressor = nn.Sequential(
nn.Dropout(dropout1),
nn.Linear(hidden, 256),
nn.ReLU(),
nn.Dropout(dropout2),
nn.Linear(256, 1),
nn.Sigmoid() # → (0, 1); we multiply by 99 at inference
)
def forward(self, input_ids, attention_mask, token_type_ids=None):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
output_attentions=True # expose attention weights for heatmap
)
cls_vec = outputs.last_hidden_state[:, 0, :] # [B, 768]
attentions = outputs.attentions # tuple of 12 layers
score = self.regressor(cls_vec).squeeze(-1) # [B]
score = score * 99.0 # scale to 0–99
return score, attentions
def get_tokenizer():
return BertTokenizer.from_pretrained("bert-base-uncased")
def load_bert_regressor(model_path: str, device: str = "cpu") -> BertRegressorModel:
"""Load a saved BertRegressorModel from disk."""
model = BertRegressorModel()
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
model.to(device)
return model
|