| |
| """Load and score an RLCD checkpoint: a ModernBERT encoder, masked mean pooling, one linear head per output, |
| softmax(logits / temperature). The same class the model card uses; train.py saves in this layout.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
|
|
|
|
| def render(messages: list[dict[str, str]]) -> str: |
| """The model's input: one `[role] content` line per message, in order.""" |
| return "\n".join(f"[{m['role']}] {m['content']}" for m in messages) |
|
|
|
|
| def pool(hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: |
| mask = attention_mask.unsqueeze(-1).to(hidden.dtype) |
| return (hidden * mask).sum(1) / mask.sum(1).clamp(min=1.0) |
|
|
|
|
| class RlcdClassifier(torch.nn.Module): |
| def __init__(self, path: Path) -> None: |
| super().__init__() |
| self.config = json.loads((path / "rlcd_config.json").read_text()) |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
| self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(path / "encoder")) |
| hidden = self.encoder.config.hidden_size |
| self.heads = torch.nn.ModuleDict( |
| {name: torch.nn.Linear(hidden, len(classes)) for name, classes in self.config["heads"].items()} |
| ) |
| |
| self.load_state_dict(load_file(str(path / "model.safetensors")), strict=True) |
| self.eval() |
|
|
| render = staticmethod(render) |
|
|
| def logits(self, texts: list[str]) -> dict[str, torch.Tensor]: |
| batch = self.tokenizer( |
| texts, padding=True, truncation=True, max_length=self.config["max_length"], return_tensors="pt" |
| ) |
| hidden = self.encoder(**batch).last_hidden_state |
| pooled = pool(hidden, batch["attention_mask"]) |
| return {name: head(pooled) for name, head in self.heads.items()} |
|
|
| @torch.no_grad() |
| def predict(self, messages: list[dict[str, str]]) -> dict[str, dict[str, float]]: |
| logits = self.logits([self.render(messages)]) |
| return { |
| name: dict(zip(self.config["heads"][name], |
| (value / self.config["temperatures"][name]).softmax(-1)[0].tolist())) |
| for name, value in logits.items() |
| } |
|
|
