"""WebOrganizer's custom gte ('new') code under transformers 5: non-persistent buffers (position_ids, RoPE cos/sin caches) come out
uninitialised (meta-device loading). Rebuild them after from_pretrained."""
import torch


def fix_buffers(m):
    for mod in m.modules():
        if hasattr(mod, "position_ids") and isinstance(getattr(mod, "position_ids"), torch.Tensor):
            n = mod.position_ids.shape[-1]; mod.position_ids = torch.arange(n, device=mod.position_ids.device).expand_as(mod.position_ids).clone()
        if hasattr(mod, "_set_cos_sin_cache"):
            dev = next((b.device for b in mod.buffers()), torch.device("cpu"))
            inv = getattr(mod, "inv_freq", None)
            if inv is not None and hasattr(mod, "base") and hasattr(mod, "dim"):
                mod.inv_freq = 1.0 / (mod.base ** (torch.arange(0, mod.dim, 2, device=dev).float() / mod.dim))
            mod._set_cos_sin_cache(seq_len=getattr(mod, "max_seq_len_cached", getattr(mod, "max_position_embeddings", 8192)), device=dev, dtype=torch.float32)
    return m


def _ext_mask(self, attention_mask, input_shape=None, device=None, dtype=None):
    """transformers 4.x ModuleUtilsMixin.get_extended_attention_mask (removed in 5): [bs, seq] -> additive [bs, 1, 1, seq] mask."""
    dtype = dtype or next(self.parameters()).dtype
    ext = attention_mask[:, None, None, :] if attention_mask.dim() == 2 else attention_mask[:, None, :, :]
    ext = ext.to(dtype=dtype)
    return (1.0 - ext) * torch.finfo(dtype).min


def patch(m):
    for mod in m.modules():
        if type(mod).__name__ == "NewModel" and not hasattr(type(mod), "get_extended_attention_mask"):
            type(mod).get_extended_attention_mask = _ext_mask
    return fix_buffers(m)