File size: 4,955 Bytes
5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 02f14f0 5c11ad4 | 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 | import os
import torch
import torch.nn as nn
from transformers import AutoModel, AutoConfig
class Model(nn.Module):
def __init__(
self,
model_dir: str,
metric_names: list[str],
dropout: float = 0.2,
hidden_dim: int | None = None
):
"""
:param model_dir: path or HF identifier for base Transformer
:param metric_names: list of your emotion columns, e.g.
['Happiness_M','Sadness_M', …,'Arousal_M']
:param dropout: dropout rate after LayerNorm
:param hidden_dim: if None, inferred from model.config.hidden_size
"""
super().__init__()
self.metric_names = metric_names
self.bert = AutoModel.from_pretrained(model_dir)
# infer hidden dim if not provided
if hidden_dim is None:
hidden_dim = self.bert.config.hidden_size
# per‐metric heads
for name in self.metric_names:
setattr(self, name, nn.Linear(hidden_dim, 1))
setattr(self, 'l_1_' + name, nn.Linear(hidden_dim, hidden_dim))
self.layer_norm = nn.LayerNorm(hidden_dim)
self.relu = nn.ReLU()
self.dropout_layer = nn.Dropout(dropout)
self.sigmoid = nn.Sigmoid()
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
# RobertaModel returns (last_hidden_state, pooled_output) when return_dict=False
_, pooled = self.bert(
input_ids = input_ids,
attention_mask = attention_mask,
return_dict = False
)
return self.rate_embedding(pooled)
def rate_embedding(self, x: torch.Tensor):
outputs = []
for name in self.metric_names:
h = getattr(self, 'l_1_' + name)(x)
# residual + norm + dropout + activation
h = self.relu(self.dropout_layer(self.layer_norm(h + x)))
out = self.sigmoid(getattr(self, name)(h))
outputs.append(out)
return outputs
def save_pretrained(self, save_directory: str):
"""
Saves:
- config.json (with custom metric_names)
- the base model files (via HF save_pretrained)
- pytorch_model.bin (full state_dict)
"""
os.makedirs(save_directory, exist_ok=True)
# extend HF config with our metric_names
config = self.bert.config
config.metric_names = self.metric_names
config.save_pretrained(save_directory)
self.bert.save_pretrained(save_directory)
torch.save(self.state_dict(), os.path.join(save_directory, 'pytorch_model.bin'))
@classmethod
def from_pretrained(
cls,
bert_model_dir: str,
state_dict_path: str | None = None,
dropout: float = 0.2,
hidden_dim: int | None = None,
metric_names: list[str] | None = None
):
"""
Loads your production model in two pieces:
1) the base Transformer from `bert_model_dir` (must be a model name or dir with a config.json)
2) your old state‐dict from `state_dict_path` (a file, e.g. no-extension)
You must pass exactly the same metric_names you trained with,
unless you previously called save_pretrained (which writes them into config.json).
"""
# 1) try to read config to recover metric_names & hidden_dim
config = AutoConfig.from_pretrained(bert_model_dir)
cfg_names = getattr(config, 'metric_names', None)
cfg_hidden = getattr(config, 'hidden_size', None)
_metric_names = metric_names or cfg_names
_hidden_dim = hidden_dim or cfg_hidden
if _metric_names is None or _hidden_dim is None:
raise ValueError(
"Must provide `metric_names` and `hidden_dim` "
"either as arguments or stored in config.json."
)
# instantiate
model = cls(
model_dir = bert_model_dir,
metric_names = _metric_names,
dropout = dropout,
hidden_dim = _hidden_dim
)
# 2) load your saved state‐dict
if state_dict_path:
if not os.path.isfile(state_dict_path):
raise FileNotFoundError(f"State‐dict file not found: {state_dict_path}")
sd = torch.load(state_dict_path, map_location='cpu')
model.load_state_dict(sd)
else:
# fallback: look for pytorch_model.bin in bert_model_dir
fallback = os.path.join(bert_model_dir, 'pytorch_model.bin')
if os.path.isfile(fallback):
sd = torch.load(fallback, map_location='cpu')
model.load_state_dict(sd)
else:
raise FileNotFoundError(
"No state‐dict provided and no pytorch_model.bin in the model_dir."
)
return model
|