Wrapper.py

from dataclasses import dataclass
from typing import List, Optional, Dict, Any, Tuple

import time
import torch
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from Segmenter import HSPPlannerConfig, Segmenter
from PageEncoder import PageEncoder
from QueryPlanner import QueryPlanner


@dataclass
class HSPWrapperConfig:
    page_size: int = 64
    flow_window: int = 4
    flash_top_k: int = 22
    anchor_pages: int = 4
    pad_token_id: Optional[int] = None
    newline_token_ids: Optional[Tuple[int, ...]] = None
    newline_token_id: int = 198
    sentence_boundary_ids: Optional[Tuple[int, ...]] = None

    allow_implicit_query: bool = False

    query_block_size: int = 64
    use_dynamic_token_weights: bool = True
    min_length_for_dynamic_weights: int = 256

    use_query_multitoken_semantic: bool = True
    min_query_tokens_for_multi: int = 4
    max_query_tokens_for_multi: int = 32
    identity_mean_weight: float = 0.7
    identity_max_weight: float = 0.3
    use_lexical_overlap_score: bool = True
    semantic_score_weight: float = 0.7
    lexical_score_weight: float = 0.3


class HSPBlackBoxWrapper:
    def __init__(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        cfg: HSPWrapperConfig,
        idf_weights: Optional[torch.Tensor] = None,
        device: Optional[torch.device] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.wrap_cfg = cfg

        if device is None:
            device = getattr(model, "device", torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.device = device
        self.model.to(self.device)

        if self.wrap_cfg.pad_token_id is None:
            if self.tokenizer.pad_token_id is not None:
                self.wrap_cfg.pad_token_id = int(self.tokenizer.pad_token_id)
            else:
                self.wrap_cfg.pad_token_id = 0
        else:
            if (
                self.tokenizer.pad_token_id is not None
                and int(self.wrap_cfg.pad_token_id) != int(self.tokenizer.pad_token_id)
            ):
                raise ValueError(
                    f"pad_token_id mismatch: wrapper cfg={self.wrap_cfg.pad_token_id}, tokenizer={self.tokenizer.pad_token_id}."
                )
        if self.wrap_cfg.newline_token_ids is None:
            auto_newline_ids: List[int] = []
            for s in ["\n", "\r\n", "\n\n"]:
                ids = self.tokenizer.encode(s, add_special_tokens=False)
                if len(ids) == 1:
                    tid = int(ids[0])
                    if tid not in auto_newline_ids:
                        auto_newline_ids.append(tid)
            if auto_newline_ids:
                self.wrap_cfg.newline_token_ids = tuple(auto_newline_ids)
                self.wrap_cfg.newline_token_id = auto_newline_ids[0]
        if self.wrap_cfg.sentence_boundary_ids is None:
            boundary_ids: List[int] = []
            if self.wrap_cfg.newline_token_ids is not None:
                boundary_ids.extend(int(x) for x in self.wrap_cfg.newline_token_ids)

            punct_candidates = ["。", "！", "？", ".", "!", "?"]
            for s in punct_candidates:
                ids = self.tokenizer.encode(s, add_special_tokens=False)
                if len(ids) == 1:
                    tid = int(ids[0])
                    if tid not in boundary_ids:
                        boundary_ids.append(tid)

            if boundary_ids:
                self.wrap_cfg.sentence_boundary_ids = tuple(boundary_ids)
        hsp_cfg = HSPPlannerConfig(
            page_size=cfg.page_size,
            flow_window=cfg.flow_window,
            flash_top_k=cfg.flash_top_k,
            anchor_pages=cfg.anchor_pages,
            pad_token_id=cfg.pad_token_id,
            newline_token_id=cfg.newline_token_id,
            newline_token_ids=cfg.newline_token_ids,
            sentence_boundary_ids=cfg.sentence_boundary_ids,
            identity_mean_weight=cfg.identity_mean_weight,
            identity_max_weight=cfg.identity_max_weight,
            lambda_semantic=cfg.semantic_score_weight,
            lambda_lexical=cfg.lexical_score_weight,
            min_query_tokens_for_multi=cfg.min_query_tokens_for_multi,
            max_query_tokens_for_multi=cfg.max_query_tokens_for_multi,
        )
        self.hsp_cfg = hsp_cfg
        boundary_ids = (
            hsp_cfg.sentence_boundary_ids
            if hsp_cfg.sentence_boundary_ids is not None
            else (
                hsp_cfg.newline_token_ids
                if (hsp_cfg.newline_token_ids is not None and len(hsp_cfg.newline_token_ids) > 0)
                else ((hsp_cfg.newline_token_id,) if hsp_cfg.newline_token_id is not None else tuple())
            )
        )
        self.segmenter = Segmenter(
            cfg=hsp_cfg,
            query_block_size=cfg.query_block_size,
            boundary_ids=boundary_ids,
            align_explicit_query_pos=False,
        )

        hidden_dim = model.config.hidden_size
        self.page_encoder = PageEncoder(hsp_cfg, hidden_dim, idf_weights=idf_weights)
        self.query_planner = QueryPlanner(hsp_cfg, query_dim=hidden_dim)

        self.to(self.device)

    def to(self, device: torch.device):
        self.device = device
        self.model.to(device)
        self.segmenter.to(device)
        self.page_encoder.to(device)
        self.query_planner.to(device)
        return self

    # Token weights (tokenizer-dependent)
    def _compute_local_token_weights(
        self,
        input_ids: torch.Tensor,          # [B, L]
        attention_mask: torch.Tensor,     # [B, L]
    ) -> torch.Tensor:
        if not self.wrap_cfg.use_dynamic_token_weights:
            return attention_mask.float()

        B, L = input_ids.shape
        pad_id = self.wrap_cfg.pad_token_id
        device = input_ids.device

        weights = torch.zeros_like(input_ids, dtype=torch.float, device=device)

        for b in range(B):
            valid = attention_mask[b].bool() if attention_mask is not None else (input_ids[b] != pad_id)
            ids_valid = input_ids[b, valid]
            L_b = ids_valid.numel()

            if L_b == 0 or L_b < self.wrap_cfg.min_length_for_dynamic_weights:
                weights[b, valid] = 1.0
                continue

            uniq, counts = ids_valid.unique(return_counts=True)
            counts = counts.float()
            w = torch.log1p(L_b / (1.0 + counts))

            w_pos = torch.zeros_like(ids_valid, dtype=torch.float, device=device)
            for uid, wv in zip(uniq.tolist(), w.tolist()):
                mask = (ids_valid == uid)
                if mask.any():
                    w_pos[mask] = wv

            maxv = float(w_pos.max().item())
            if maxv > 0:
                w_pos = w_pos / maxv

            weights[b, valid] = w_pos

        weights = weights + 1e-6
        return weights

    # Build tokenized inputs (tokenizer-dependent)
    def _build_inputs_from_texts(
        self,
        contexts: List[str],
        questions: Optional[List[str]],
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        B = len(contexts)
        if questions is None:
            questions = [""] * B
        else:
            assert len(contexts) == len(questions)

        ctx_enc = [self.tokenizer.encode(c, add_special_tokens=False) for c in contexts]

        implicit_mode = False
        if self.wrap_cfg.allow_implicit_query:
            no_query_flags = [((q is None) or (str(q).strip() == "")) for q in questions]
            if all(no_query_flags):
                implicit_mode = True

        if implicit_mode:
            q_enc = [[] for _ in range(B)]
        else:
            q_enc = [self.tokenizer.encode((q if q is not None else ""), add_special_tokens=False) for q in questions]

        input_ids_list = []
        explicit_qp_list = []

        for i in range(B):
            ids_ctx = ctx_enc[i]
            ids_q = q_enc[i]
            all_ids = ids_ctx if implicit_mode else (ids_ctx + ids_q)
            if len(all_ids) == 0:
                all_ids = [self.wrap_cfg.pad_token_id]
                qp = 0
            else:
                qp = len(ids_ctx)
            input_ids_list.append(all_ids)
            if not implicit_mode:
                explicit_qp_list.append(qp)

        max_len = max(len(x) for x in input_ids_list)
        pad_id = self.wrap_cfg.pad_token_id

        input_ids = torch.full((B, max_len), pad_id, dtype=torch.long, device=self.device)
        attention_mask = torch.zeros((B, max_len), dtype=torch.long, device=self.device)

        for i in range(B):
            L_i = len(input_ids_list[i])
            input_ids[i, :L_i] = torch.tensor(input_ids_list[i], dtype=torch.long, device=self.device)
            attention_mask[i, :L_i] = 1

        explicit_qp = None if implicit_mode else torch.tensor(explicit_qp_list, dtype=torch.long, device=self.device)
        return input_ids, attention_mask, explicit_qp

    # Compress inputs for prefill
    @torch.no_grad()
    def compress_inputs_for_prefill(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        explicit_query_pos: Optional[torch.Tensor],
    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
        B, L = input_ids.shape
        aligned_qp, split_info, layout = self.segmenter.build_layout(
            input_ids=input_ids,
            attention_mask=attention_mask,
            explicit_query_pos=explicit_query_pos,
        )
        if hasattr(self.model, "get_input_embeddings") and self.model.get_input_embeddings() is not None:
            embed = self.model.get_input_embeddings()
            hidden = embed(input_ids)
        elif hasattr(self.model, "model") and hasattr(self.model.model, "embed_tokens"):
            hidden = self.model.model.embed_tokens(input_ids)
        elif hasattr(self.model, "transformer") and hasattr(self.model.transformer, "wte"):
            hidden = self.model.transformer.wte(input_ids)
        else:
            raise RuntimeError("Cannot locate input embedding layer for wrapped model.")
        token_level_weights = self._compute_local_token_weights(input_ids, attention_mask)
        block_repr = self.page_encoder(
            hidden_states=hidden,
            layout=layout,
            input_ids=input_ids,
            token_level_weights=token_level_weights,
        )
        token_valid = layout.token_valid
        query_hidden_list = []
        pos_idx = torch.arange(L, device=self.device)

        w_mean = getattr(self.hsp_cfg, "identity_mean_weight", 0.7)
        w_max = getattr(self.hsp_cfg, "identity_max_weight", 0.3)
        w_sum = w_mean + w_max
        if w_sum <= 0:
            w_mean_eff, w_max_eff = 1.0, 0.0
        else:
            w_mean_eff, w_max_eff = w_mean / w_sum, w_max / w_sum

        has_idf = getattr(self.page_encoder, "idf_weights", None) is not None

        for b in range(B):
            split_b = split_info[b]
            qs, qe = int(split_b.query_start), int(split_b.query_end)
            if qe < qs:
                qp = int(aligned_qp[b].item())
                query_hidden_list.append(hidden[b, qp])
                continue
            span_mask = (pos_idx >= qs) & (pos_idx <= qe) & token_valid[b]
            if not span_mask.any():
                qp = int(aligned_qp[b].item())
                query_hidden_list.append(hidden[b, qp])
                continue
            h_b = hidden[b]
            ids_b = input_ids[b]
            if has_idf and input_ids is not None:
                idf_vec = self.page_encoder.idf_weights
                idf_b = idf_vec[ids_b]
                weights = idf_b * span_mask.float()
                w_sum_b = float(weights.sum().item())
                if w_sum_b > 1e-6:
                    mean_b = (h_b * weights.unsqueeze(-1)).sum(dim=0) / w_sum_b
                else:
                    mean_b = h_b[span_mask].mean(dim=0)
            else:
                mean_b = h_b[span_mask].mean(dim=0)
            max_b, _ = h_b[span_mask].max(dim=0)
            pooled_b = w_mean_eff * mean_b + w_max_eff * max_b
            query_hidden_list.append(pooled_b)

        query_hidden = torch.stack(query_hidden_list, dim=0)
        query_token_hidden_list = []
        query_token_weight_list = []

        min_q = getattr(self.hsp_cfg, "min_query_tokens_for_multi", 4)
        max_q = getattr(self.hsp_cfg, "max_query_tokens_for_multi", 32)
        use_multi = self.wrap_cfg.use_query_multitoken_semantic and min_q > 0

        for b in range(B):
            split_b = split_info[b]
            qs, qe = int(split_b.query_start), int(split_b.query_end)
            if qe < qs or not use_multi:
                query_token_hidden_list.append(None)
                query_token_weight_list.append(None)
                continue
            span_mask = (pos_idx >= qs) & (pos_idx <= qe) & token_valid[b]
            idx_span = pos_idx[span_mask]
            Tq = idx_span.numel()
            if Tq < min_q:
                query_token_hidden_list.append(None)
                query_token_weight_list.append(None)
                continue
            h_b = hidden[b]
            w_b = token_level_weights[b]
            w_span = w_b[idx_span]
            top_k = min(Tq, max_q)
            vals, indices = torch.topk(w_span, k=top_k, largest=True, sorted=True)
            idx_top = idx_span[indices]
            h_top = h_b[idx_top]
            w_top = vals / (vals.sum() + 1e-6)
            query_token_hidden_list.append(h_top)
            query_token_weight_list.append(w_top)
        keep_pages = self.query_planner(
            block_repr=block_repr,
            layout=layout,
            query_hidden=query_hidden,
            query_pos=aligned_qp,
            input_ids=input_ids,
            token_level_weights=token_level_weights,
            split_results=split_info,
            query_token_hidden_list=query_token_hidden_list,
            query_token_weight_list=query_token_weight_list,
        )
        token2page = layout.token2page               # [B, L]
        token_valid = layout.token_valid             # [B, L] bool
        keep_pages_bool = keep_pages.bool()          # [B, P] bool
        token2page_clamped = token2page.clamp(min=0) # avoid -1 gather; token_valid will mask anyway
        keep_token = token_valid & keep_pages_bool.gather(1, token2page_clamped)
        boundary_ids = (
            self.hsp_cfg.sentence_boundary_ids
            if self.hsp_cfg.sentence_boundary_ids is not None
            else (
                self.hsp_cfg.newline_token_ids
                if (self.hsp_cfg.newline_token_ids is not None and len(self.hsp_cfg.newline_token_ids) > 0)
                else ((self.hsp_cfg.newline_token_id,) if self.hsp_cfg.newline_token_id is not None else tuple())
            )
        )
        kept_context_token_indices: List[List[int]] = []
        context_lens: List[int] = []
        for b in range(B):
            qp = int(aligned_qp[b].item())
            qp = max(0, min(qp, L))
            context_lens.append(qp)
            if qp <= 0:
                kept_context_token_indices.append([])
                continue
            keep_ctx = keep_token[b, :qp].clone()
            valid_ctx = token_valid[b, :qp]
            if boundary_ids is not None and len(boundary_ids) > 0:
                ids_slice = input_ids[b, :qp]
                is_boundary = torch.zeros(qp, dtype=torch.bool, device=input_ids.device)
                for bid in boundary_ids:
                    is_boundary |= (ids_slice == int(bid))
                boundary_pos = torch.nonzero(is_boundary, as_tuple=False).flatten().tolist()
                start = 0
                for p in boundary_pos:
                    end = int(p) + 1
                    if end <= start:
                        continue
                    if keep_ctx[start:end].any():
                        keep_ctx[start:end] = valid_ctx[start:end]
                    start = end
                if start < qp and keep_ctx[start:qp].any():
                    keep_ctx[start:qp] = valid_ctx[start:qp]
            kept_idx = torch.nonzero(keep_ctx & valid_ctx, as_tuple=False).flatten().detach().cpu().tolist()
            kept_context_token_indices.append([int(x) for x in kept_idx])
        compressed = self.segmenter.compress(
            input_ids=input_ids,
            attention_mask=attention_mask,
            layout=layout,
            keep_pages=keep_pages,
            query_pos=aligned_qp,
        )
        L_comp = compressed["input_ids"].size(1)
        stats = {
            "aligned_query_pos": [int(x) for x in aligned_qp.detach().cpu().tolist()],
            "context_len": context_lens,
            "kept_context_token_indices": kept_context_token_indices,
            "original_len": int(L),
            "compressed_len": int(L_comp),
            "compression_ratio": float(L_comp / max(L, 1)),
        }
        return compressed, stats

    # Batch generation (dense or BEAVER)
    @torch.no_grad()
    def generate_batch(
        self,
        contexts: List[str],
        questions: List[str],
        max_new_tokens: int = 128,
        use_hsp: bool = True,
        **gen_kwargs,
    ) -> Dict[str, Any]:
        assert len(contexts) == len(questions)
        self.model.eval()

        input_ids, attention_mask, explicit_qp = self._build_inputs_from_texts(contexts, questions)
        prompt_lens = attention_mask.sum(dim=1).tolist()

        if not use_hsp:
            t0 = time.time()
            outputs = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                use_cache=True,
                **gen_kwargs,
            )
            t1 = time.time()

            texts: List[str] = []
            for i in range(outputs.size(0)):
                pl = int(prompt_lens[i]) if i < len(prompt_lens) else 0
                answer_tokens = outputs[i, pl:]
                texts.append(self.tokenizer.decode(answer_tokens, skip_special_tokens=True))

            meta = {
                "mode": "dense",
                "time_total": float(t1 - t0),
                "original_len": int(input_ids.size(1)),
                "compressed_len": int(input_ids.size(1)),
                "compression_ratio": 1.0,
            }
            return {"outputs": texts, "meta": meta}
        t0 = time.time()
        compressed_inputs, stats = self.compress_inputs_for_prefill(
            input_ids=input_ids,
            attention_mask=attention_mask,
            explicit_query_pos=explicit_qp,
        )
        t1 = time.time()

        comp_attn = compressed_inputs["attention_mask"]
        comp_prompt_lens = comp_attn.sum(dim=1).tolist()
        gen_outputs = self.model.generate(
            input_ids=compressed_inputs["input_ids"],
            attention_mask=compressed_inputs["attention_mask"],
            max_new_tokens=max_new_tokens,
            use_cache=True,
            **gen_kwargs,
        )
        t2 = time.time()

        texts: List[str] = []
        for i in range(gen_outputs.size(0)):
            pl = int(comp_prompt_lens[i]) if i < len(comp_prompt_lens) else 0
            answer_tokens = gen_outputs[i, pl:]
            texts.append(self.tokenizer.decode(answer_tokens, skip_special_tokens=True))

        meta = {
            "mode": "hsp",
            "time_prefill": float(t1 - t0),
            "time_total": float(t2 - t0),
            "original_len": stats["original_len"],
            "compressed_len": stats["compressed_len"],
            "compression_ratio": stats["compression_ratio"],
        }
        return {"outputs": texts, "meta": meta}