| """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 |
