File size: 2,463 Bytes
2778af7 | 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 | """Cross-encoder specialist: state + question + one option in a single encoder pass."""
from __future__ import annotations
import json
import os
import torch
import torch.nn as nn
from safetensors.torch import load_file, save_file
from transformers import AutoConfig, AutoModel
class PlutoCrossEncoder(nn.Module):
def __init__(self, backbone_name: str | None = None, backbone_config=None, vocab_size: int | None = None):
super().__init__()
if backbone_config is None:
self.backbone = AutoModel.from_pretrained(backbone_name, attn_implementation="sdpa")
else:
self.backbone = AutoModel.from_config(backbone_config, attn_implementation="sdpa")
if vocab_size is not None:
self.backbone.resize_token_embeddings(vocab_size)
h = self.backbone.config.hidden_size
self.head = nn.Sequential(
nn.Linear(h, h),
nn.GELU(),
nn.Dropout(0.1),
nn.LayerNorm(h),
nn.Linear(h, 1),
)
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
return self.head(h).squeeze(-1).float()
def save_cross_encoder(model: PlutoCrossEncoder, tokenizer, out_dir: str, extra: dict) -> None:
os.makedirs(out_dir, exist_ok=True)
save_file({k: v.detach().cpu() for k, v in model.state_dict().items()}, os.path.join(out_dir, "model.safetensors"))
tokenizer.save_pretrained(out_dir)
model.backbone.config.to_json_file(os.path.join(out_dir, "backbone_config.json"))
cfg = {
"format": "pluto-ce-1",
"vocab_size": model.backbone.config.vocab_size,
**extra,
}
with open(os.path.join(out_dir, "pluto_ce_config.json"), "w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
def load_cross_encoder(path: str, device: str | torch.device = "cpu") -> tuple[PlutoCrossEncoder, dict]:
cfg = json.load(open(os.path.join(path, "pluto_ce_config.json"), encoding="utf-8"))
backbone_cfg = AutoConfig.for_model(
**json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8"))
)
model = PlutoCrossEncoder(backbone_config=backbone_cfg, vocab_size=cfg["vocab_size"])
model.load_state_dict(load_file(os.path.join(path, "model.safetensors")), strict=True)
model.to(device).eval()
return model, cfg
|