File size: 1,323 Bytes
de23f96 | 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 | """Load the query-aware snippet extraction checkpoint from a local directory."""
from pathlib import Path
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
class SentenceCompressor(nn.Module):
"""ModernBERT encoder with a token-level keep/drop classification head."""
def __init__(self, base: str, dropout: float = 0.1):
super().__init__()
self.encoder = AutoModel.from_pretrained(base, attn_implementation="sdpa")
self.dropout = nn.Dropout(dropout)
self.head = nn.Linear(self.encoder.config.hidden_size, 1)
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
output = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
return self.head(self.dropout(output.last_hidden_state)).squeeze(-1)
def load_model(model_dir: str | Path, device: str = "cpu"):
"""Return the fine-tuned model and tokenizer in evaluation mode."""
model_dir = Path(model_dir)
checkpoint = torch.load(model_dir / "model.pt", map_location=device, weights_only=False)
model = SentenceCompressor(base=checkpoint["args"]["base"]).to(device)
model.load_state_dict(checkpoint["model"])
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
return model, tokenizer |