File size: 3,158 Bytes
6b77899 |
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 |
"""
Utility module for loading and running the finetuned ModernBERT reward model.
The model mirrors the architecture defined in `mosaic_bert_training.py`:
- base encoder: answerdotai/ModernBERT-base (8k context support)
- pooling: attention-mask-weighted mean pooling
- head: single linear layer + sigmoid to output a score in [0, 1]
"""
import os
from typing import Optional
import torch
from transformers import AutoModel
class BERTRewardModel(torch.nn.Module):
"""ModernBERT encoder with a sigmoid regression head."""
def __init__(self, model_name: str = "answerdotai/ModernBERT-base"):
super().__init__()
self.bert = AutoModel.from_pretrained(
model_name,
reference_compile=False,
attn_implementation="eager",
)
self.classifier = torch.nn.Linear(self.bert.config.hidden_size, 1)
self.sigmoid = torch.nn.Sigmoid()
def forward(self, input_ids, attention_mask, labels: Optional[torch.Tensor] = None):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
last_hidden_state = outputs.last_hidden_state
attention_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
sum_hidden = torch.sum(last_hidden_state * attention_mask_expanded, dim=1)
sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
pooled_output = sum_hidden / sum_mask
logits = self.classifier(pooled_output)
scores = self.sigmoid(logits).squeeze(-1)
loss = None
if labels is not None:
loss_fct = torch.nn.MSELoss()
loss = loss_fct(scores, labels) * 100
return {"loss": loss, "scores": scores, "logits": logits}
def load_finetuned_model(model_dir: str, device: Optional[str] = None) -> BERTRewardModel:
"""
Load the finetuned ModernBERT reward model from `model_dir`.
Args:
model_dir: Path containing model.safetensors (preferred) or pytorch_model.bin.
device: Optional torch device string. Defaults to CUDA if available else CPU.
"""
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_device = torch.device(device)
model = BERTRewardModel()
model.to(torch_device)
state_dict = None
safetensors_path = os.path.join(model_dir, "model.safetensors")
bin_path = os.path.join(model_dir, "pytorch_model.bin")
if os.path.exists(safetensors_path):
try:
from safetensors.torch import load_file
state_dict = load_file(safetensors_path)
except ImportError as exc:
print(f"⚠️ safetensors not available ({exc}); falling back to pytorch_model.bin if present.")
if state_dict is None and os.path.exists(bin_path):
state_dict = torch.load(bin_path, map_location=torch_device)
if state_dict is None:
raise FileNotFoundError(
f"Could not find model weights in {model_dir}. "
"Expected either model.safetensors or pytorch_model.bin."
)
model.load_state_dict(state_dict)
model.eval()
return model
|