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