Upload 4 files

utils/complete_model.py

@@ -0,0 +1,276 @@
+import os
+import torch
+import torch.nn as nn
+from transformers import AutoModel, GPT2Tokenizer
+
+from utils.modifiedGPT2 import create_decoder
+
+from utils.layer_mask import gaussian_layer_stack_pipeline
+
+
+class DINOEncoder(nn.Module):
+    def __init__(self, model_id="facebook/dinov3-vits16-pretrain-lvd1689m", freeze=True):
+        super().__init__()
+        self.model = AutoModel.from_pretrained(model_id)
+        if freeze:
+            for p in self.model.parameters():
+                p.requires_grad = False
+    @torch.no_grad()
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        pixel_values: [B, C, H, W]
+        returns patches: [B, Np, Cenc]
+        """
+        out = self.model(pixel_values=pixel_values)
+        tokens = out.last_hidden_state  # [B, 1+Np, Cenc] (CLS + patches) for ViT-like
+        # Skip a few special tokens if your backbone adds them; adjust as needed.
+        patches = tokens[:, 5:, :]  # [B, Np, Cenc]
+        return patches
+
+class DinoUNet(nn.Module):
+    def __init__(self, model_name="facebook/dinov3-convnext-small-pretrain-lvd1689m", freeze=True):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name)
+        # NOTE: confirm channels of the chosen hidden state; 768 is common for small convnext/dinov3
+        self.channel_adapter = nn.Conv2d(768, 512, kernel_size=1)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(512, 256, 3, padding=1), nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(256, 128, 2, stride=2), nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(128, 64, 2, stride=2), nn.ReLU(inplace=True),
+            nn.Conv2d(64, 1, 1)
+        )
+        if freeze:
+            for m in (self.encoder, self.channel_adapter, self.decoder):
+                for p in m.parameters():
+                    p.requires_grad = False
+    
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, num_layers: int) -> torch.Tensor:
+        """
+        x: [B, C, H, W]; returns mask: [B, 1, H', W'] (your upsampling stack defines H',W')
+        """
+        enc_feats = self.encoder(x, output_hidden_states=True, return_dict=True)
+        # take the last 4D feature map from hidden_states
+        feats = next(h for h in reversed(enc_feats.hidden_states) if isinstance(h, torch.Tensor) and h.ndim == 4)
+        feats = self.channel_adapter(feats)
+        pred = self.decoder(feats)                    # (B,1,h,w)
+        _, _, segmentation_mask = gaussian_layer_stack_pipeline(pred, n_layers = num_layers)
+        return segmentation_mask    # [B, num_layers, h, w]
+
+
+class LinearProjection(nn.Module):
+    def __init__(self, input_dim=384, output_dim=768, freeze=False):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, output_dim)
+        if freeze:
+            for p in self.proj.parameters():
+                p.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, Np, input_dim] -> [B, Np, output_dim]
+        return self.proj(x)
+
+
+class CustomModel(nn.Module):
+    def __init__(
+        self,
+        device: str = "cuda",
+        ENCODER_MODEL_PATH: str | None = "dino_encoder.pth",
+        SEGMENTER_MODEL_PATH: str | None = "dino_segmenter.pth",
+        DECODER_MODEL_PATH: str | None = "dino_decoder.pth",
+        LINEAR_PROJECTION_PATH: str | None = "linear_projection.pth",
+        freeze_encoder: bool = True,
+        freeze_segmenter: bool = True,
+        freeze_linear_projection: bool = False,
+        freeze_decoder: bool = False,
+        attention_implementation: str = "sdpa",
+    ):
+        super().__init__()
+        self.device = torch.device(device)
+
+        # Encoder
+        self.encoder = DINOEncoder()
+        if ENCODER_MODEL_PATH and os.path.exists(ENCODER_MODEL_PATH):
+            self.encoder.load_state_dict(torch.load(ENCODER_MODEL_PATH, map_location="cpu"), strict=False)
+            print("Loaded encoder weights from", ENCODER_MODEL_PATH)
+        if freeze_encoder:
+            self.encoder.eval()
+
+        # Segmenter
+        self.segmenter = DinoUNet()
+        if SEGMENTER_MODEL_PATH and os.path.exists(SEGMENTER_MODEL_PATH):
+            self.segmenter.load_state_dict(torch.load(SEGMENTER_MODEL_PATH, map_location="cpu"), strict=False)
+            print("Loaded segmenter weights from", SEGMENTER_MODEL_PATH)
+        if freeze_segmenter:
+            self.segmenter.eval()
+
+        # Decoder (modified GPT-2)
+        self.decoder = create_decoder(attention=attention_implementation)  # must expose .config.hidden_size & .config.num_hidden_layers
+        if DECODER_MODEL_PATH and os.path.exists(DECODER_MODEL_PATH):
+            self.decoder.load_state_dict(torch.load(DECODER_MODEL_PATH, map_location="cpu"), strict=False)
+            print("Loaded decoder weights from", DECODER_MODEL_PATH)
+        if freeze_decoder:
+            self.decoder.eval()
+
+        # Linear projection: DINO hidden -> GPT2 hidden
+        enc_h = self.encoder.model.config.hidden_size
+        dec_h = self.decoder.config.hidden_size
+        self.linear_projection = LinearProjection(input_dim=enc_h, output_dim=dec_h)
+        if LINEAR_PROJECTION_PATH and os.path.exists(LINEAR_PROJECTION_PATH):
+            self.linear_projection.load_state_dict(torch.load(LINEAR_PROJECTION_PATH, map_location="cpu"), strict=False)
+            print("Loaded linear projection weights from", LINEAR_PROJECTION_PATH)
+        if freeze_linear_projection:
+            self.linear_projection.eval()
+
+        # Tokenizer (pad token for GPT-2)
+        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.pad_token_id = self.tokenizer.pad_token_id  # ✅ use ID, not string
+
+        self.num_layers = self.decoder.config.num_hidden_layers
+
+        # move everything once
+        self.to(self.device)
+
+    def forward(self, pixel_values: torch.Tensor, tgt_ids: torch.Tensor | None = None, **kwargs) -> dict:
+        """
+        pixel_values: [B,C,H,W], float
+        tgt_ids: [B,T], long (token IDs), padded with pad_token_id if any padding is present
+        """
+        pixel_values = pixel_values.to(self.device, non_blocking=True)
+
+        # Visual path
+        patches = self.encoder(pixel_values)                           # [B,Np,Cenc]
+        projected_patches = self.linear_projection(patches)            # [B,Np,n_embd]
+
+        # Segmentation path per layer
+        segmented_layers = self.segmenter(pixel_values, self.num_layers) # [B,n_layers,H,W] (per current decoder)
+
+        # Text path (optional teacher-forced training)
+        labels = None
+        if tgt_ids is not None:
+            if tgt_ids.dtype != torch.long:
+                tgt_ids = tgt_ids.long()
+            tgt_ids = tgt_ids.to(self.device, non_blocking=True)       # [B,T]
+            text_embeds = self.decoder.transformer.wte(tgt_ids)        # [B,T,n_embd]
+            inputs_embeds = torch.cat([projected_patches, text_embeds], dim=1)  # [B,Np+T,n_embd]
+
+            # Labels: ignore prefix tokens (vision) and PADs in text
+            B, Np, _ = projected_patches.shape
+            labels_prefix = torch.full((B, Np), -100, device=self.device, dtype=torch.long)
+            text_labels = tgt_ids.clone()
+            text_labels[text_labels == self.pad_token_id] = -100       # ✅ compare to ID
+            labels = torch.cat([labels_prefix, text_labels], dim=1)    # [B,Np+T]
+        else:
+            inputs_embeds = projected_patches
+
+        # Decoder forward
+        out = self.decoder(inputs_embeds=inputs_embeds, segmentation_mask=segmented_layers, labels=labels, **kwargs)
+        return out
+    
+    @torch.inference_mode()
+    def generate(
+        self,
+        pixel_values: torch.Tensor,
+        max_new_tokens: int = 100,
+        output_attentions: bool = False,
+    ) -> torch.Tensor:
+        """
+        pixel_values: [B,C,H,W], float
+        returns generated_ids: [B, T]
+        """
+        pixel_values = pixel_values.to(self.device, non_blocking=True)
+
+        # Visual path
+        patches = self.encoder(pixel_values)                           # [B,Np,Cenc]
+        projected_patches = self.linear_projection(patches)            # [B,Np,n_embd]
+
+        # Segmentation path per layer
+        segmented_layers = self.segmenter(pixel_values, self.num_layers) # [B,n_layers,H,W] (per current decoder)
+
+        # Generate
+        output = self.decoder.generate(
+            inputs_embeds=projected_patches,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            repetition_penalty=1.2,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.pad_token_id,
+            use_cache=True,
+            segmentation_mask=segmented_layers,
+            prefix_allowed_length=0,
+            plot_attention_mask=False,
+            plot_attention_mask_layer=[],
+            plot_attention_map=False,
+            plot_attention_map_layer=[],
+            plot_attention_map_generation=0,
+            output_attentions=output_attentions,
+            return_dict_in_generate=True,
+        )
+        # Remove prefix tokens (vision)
+        generated_ids = output.sequences#[:, projected_patches.shape[1]:]   # [B,T]
+        generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+        return generated_ids, generated_text, output.attentions if output_attentions else None
+
+def create_complete_model(device: str = "cuda", **kwargs) -> CustomModel:
+    model = CustomModel(device=device, **kwargs)
+    return model
+
+def save_complete_model(model: CustomModel, save_path: str, device: str = "cuda") -> None:
+    # Ensure folder exists
+    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+
+    # Save on CPU to keep checkpoint portable
+    orig_device = next(model.parameters()).device
+    model.to("cpu")
+    torch.save(model.state_dict(), save_path)
+    print(f"Saved complete model weights to {save_path}")
+
+    # Restore model device
+    model.to(device if isinstance(device, str) else orig_device)
+
+def save_checkpoint(model: CustomModel, optimizer: torch.optim.Optimizer, save_path: str) -> None:
+    # Ensure folder exists
+    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+
+    checkpoint = {
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+    }
+    torch.save(checkpoint, save_path)
+    print(f"Saved checkpoint to {save_path}")
+
+def load_complete_model(model: CustomModel, load_path: str, device: str = "cpu", strict: bool = True) -> CustomModel:
+    if not os.path.exists(load_path):
+        print(f"No weights found at {load_path}")
+        model.to(device)
+        return model
+
+    # Load to CPU first, then move to target device
+    state = torch.load(load_path, map_location="cpu")
+    missing, unexpected = model.load_state_dict(state, strict=strict)
+    if not strict:
+        if missing:
+            print(f"[load warning] Missing keys: {missing}")
+        if unexpected:
+            print(f"[load warning] Unexpected keys: {unexpected}")
+
+    model.to(device)
+    print(f"Loaded complete model weights from {load_path}")
+    return model
+
+def load_checkpoint(model: CustomModel, optimizer: torch.optim.Optimizer, load_path: str, device: str = "cpu") -> tuple[CustomModel, torch.optim.Optimizer]:
+    if not os.path.exists(load_path):
+        print(f"No checkpoint found at {load_path}")
+        model.to(device)
+        return model, optimizer
+
+    # Load to CPU first, then move to target device
+    checkpoint = torch.load(load_path, map_location="cpu")
+    model.load_state_dict(checkpoint["model_state_dict"])
+    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+
+    model.to(device)
+    print(f"Loaded checkpoint from {load_path}")
+    return model, optimizer

utils/layer_mask.py

@@ -0,0 +1,224 @@
+import torch
+import torch.nn.functional as F
+import math
+import numpy as np
+import matplotlib.pyplot as plt
+
+@torch.no_grad()
+def gaussian_layer_stack_pipeline(
+    x: torch.Tensor,
+    n_layers: int,
+    base_ksize: int = 3,
+    ksize_growth: int = 2,
+    sigma: float | None = None,
+    eps: float = 1e-8,
+):
+    """
+    All-in-one GPU batch pipeline:
+      1) Per-sample min-max normalize to [0,1]
+      2) Resize to (32,32)
+      3) Apply L Gaussian blurs with increasing kernel size in a single
+         horizontal conv + single vertical conv using depthwise groups
+         (via a shared max kernel padded with zeros)
+      4) Renormalize each layer to [0,1]
+      5) Return stacked (B,L,32,32), flat (B,L,1024), tiled (B,L,1024,1024 view)
+
+    Args:
+      x: (B,H,W) or (B,1,H,W) tensor (any device/dtype)
+      n_layers: number of layers
+      base_ksize: starting odd kernel size (e.g., 3)
+      ksize_growth: increment per layer (e.g., 2) -> ensures odd sizes
+      sigma: if None, uses (ksize-1)/6 per layer; else fixed sigma for all
+      eps: small number for safe division
+
+    Returns:
+      stacked: (B, n_layers, 32, 32)  float on x.device
+      flat:    (B, n_layers, 1024)
+      tiled:   (B, n_layers, 1024, 1024)  (expand view; memory-cheap)
+    """
+    assert n_layers >= 1, "n_layers must be >= 1"
+
+    # ---- Ensure 4D, 1 channel; cast to float (stay on same device) ----
+    if x.ndim == 3:
+        x = x.unsqueeze(1)  # (B,1,H,W)
+    elif x.ndim != 4 or x.shape[1] not in (1,):
+        raise ValueError(f"Expected (B,H,W) or (B,1,H,W); got {tuple(x.shape)}")
+    x = x.float()
+
+    B, _, H, W = x.shape
+
+    # ---- Per-sample min-max normalize to [0,1] ----
+    xmin = x.amin(dim=(2, 3), keepdim=True)
+    xmax = x.amax(dim=(2, 3), keepdim=True)
+    denom = (xmax - xmin).clamp_min(eps)
+    x = (x - xmin) / denom  # (B,1,H,W) in [0,1]
+
+    # ---- Resize to 32x32 on GPU ----
+    x = F.interpolate(x, size=(32, 32), mode="bilinear", align_corners=False)  # (B,1,32,32)
+
+    # ---- Prepare per-layer kernel sizes (odd) ----
+    ksizes = []
+    for i in range(n_layers, 0, -1):  # to keep your original ordering: L...1
+        k = base_ksize + i * ksize_growth
+        k = int(k)
+        if k % 2 == 0:
+            k += 1
+        k = max(k, 1)
+        ksizes.append(k)
+
+    Kmax = max(ksizes)
+    pad = Kmax // 2
+
+    # ---- Build per-layer 1D Gaussian vectors and embed into shared Kmax kernel ----
+    # We create horizontal weights of shape (L,1,1,Kmax) and vertical (L,1,Kmax,1)
+    device, dtype = x.device, x.dtype
+    weight_h = torch.zeros((n_layers, 1, 1, Kmax), device=device, dtype=dtype)
+    weight_v = torch.zeros((n_layers, 1, Kmax, 1), device=device, dtype=dtype)
+
+    for idx, k in enumerate(ksizes):
+        # choose sigma
+        sig = sigma if (sigma is not None and sigma > 0) else (k - 1) / 6.0
+        r = k // 2
+        xp = torch.arange(-r, r + 1, device=device, dtype=dtype)
+        g = torch.exp(-(xp * xp) / (2.0 * sig * sig))
+        g = g / g.sum()  # (k,)
+
+        # center g into Kmax with zeros around
+        start = (Kmax - k) // 2
+        end = start + k
+
+        # horizontal row
+        weight_h[idx, 0, 0, start:end] = g  # (1 x Kmax)
+
+        # vertical column
+        weight_v[idx, 0, start:end, 0] = g  # (Kmax x 1)
+
+    # ---- Duplicate input across L channels (depthwise groups) ----
+    xL = x.expand(B, n_layers, 32, 32).contiguous()  # (B,L,32,32)
+
+    # ---- Separable Gaussian blur with a single pass per axis (groups=L) ----
+    # Horizontal
+    xh = F.pad(xL, (pad, pad, 0, 0), mode="reflect")
+    xh = F.conv2d(xh, weight=weight_h, bias=None, stride=1, padding=0, groups=n_layers)  # (B,L,32,32)
+
+    # Vertical
+    xv = F.pad(xh, (0, 0, pad, pad), mode="reflect")
+    yL = F.conv2d(xv, weight=weight_v, bias=None, stride=1, padding=0, groups=n_layers)  # (B,L,32,32)
+
+    # ---- Renormalize each layer to [0,1] (per-sample, per-layer) ----
+    y_min = yL.amin(dim=(2, 3), keepdim=True)
+    y_max = yL.amax(dim=(2, 3), keepdim=True)
+    y_den = (y_max - y_min).clamp_min(eps)
+    stacked = (yL - y_min) / y_den  # (B,L,32,32) in [0,1]
+
+    # ---- Flatten + tile (expand view; caution w/ later materialization) ----
+    flat = stacked.reshape(B, n_layers, 32 * 32)               # (B,L,1024)
+    tiled = flat.unsqueeze(-2).expand(-1, -1, 32 * 32, -1)     # (B,L,1024,1024) view
+
+    return stacked, flat, tiled
+
+def plot_layers_any(
+    x,
+    *,
+    max_batches=None,
+    vlim=(0, 1),
+    one_indexed: bool = False,
+    max_cols: int = 6,
+):
+    """
+    Plot layers for each batch sample in separate figures.
+
+    Accepts:
+      - stacked: (B, L, H, W)
+      - flat:    (B, L, HW)
+      - tiled:   (B, L, HW, HW)
+
+    Behavior:
+      - Creates one figure PER BATCH (up to `max_batches`).
+      - At most `max_cols` layers per row (default 6).
+      - Column headers: 'Layer {i}' descending from n-1 -> 0 (or n -> 1 if one_indexed=True).
+      - Figure title per batch: 'Masks for input {i} out of {B}'.
+
+    Returns:
+      A list of (fig, axes) tuples, one per plotted batch.
+    """
+    # ---- Normalize input to torch ----
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor or np.ndarray, got {type(x)}")
+
+    if x.ndim not in (3, 4):
+        raise ValueError(f"Expected ndim 3 or 4, got shape {tuple(x.shape)}")
+
+    # ---- Convert to (B, L, H, W) 'stacked' ----
+    if x.ndim == 4:
+        B, L, A, B_ = x.shape
+        if A == B_:
+            # Could be stacked (H==W) or tiled (HW x HW). Heuristic: if A is a perfect square
+            # and reasonably large (e.g., 1024), treat as tiled and collapse to flat.
+            s = int(math.isqrt(A))
+            if s * s == A and A >= 64:
+                flat = x[..., 0, :].detach()  # (B, L, HW)
+                H = W = s
+                stacked = flat.reshape(B, L, H, W)
+            else:
+                stacked = x.detach()
+        else:
+            stacked = x.detach()
+    else:
+        # x.ndim == 3 -> (B, L, HW)
+        B, L, HW = x.shape
+        s = int(math.isqrt(HW))
+        if s * s != HW:
+            if HW != 32 * 32:
+                raise ValueError(
+                    f"Cannot infer square image size from HW={HW}. "
+                    f"Provide stacked (B,L,H,W) or flat with square HW."
+                )
+            s = 32
+        H = W = s
+        stacked = x.detach().reshape(B, L, H, W)
+
+    # Ensure float & CPU for plotting
+    stacked = stacked.to(torch.float32).cpu().numpy()
+
+    # ---- Batch selection ----
+    B, L, H, W = stacked.shape
+    plot_B = B if max_batches is None else max(1, min(B, int(max_batches)))
+
+    # ---- Layout params ----
+    cols = max(1, int(max_cols))
+    rows_needed = lambda L: (L + cols - 1) // cols
+
+    figs = []
+    for b in range(plot_B):
+        # number of rows for this batch
+        r = rows_needed(L)
+        fig, axes = plt.subplots(r, cols, figsize=(cols * 3, r * 3), squeeze=False)
+        fig.suptitle(f"Masks for input {b} out of {B}", fontsize=12, y=1.02)
+
+        for l in range(L):
+            rr = l // cols
+            cc = l % cols
+            ax = axes[rr, cc]
+            if vlim is None:
+                ax.imshow(stacked[b, l], cmap="gray")
+            else:
+                ax.imshow(stacked[b, l], cmap="gray", vmin=vlim[0], vmax=vlim[1])
+            ax.axis("off")
+
+            # Set column titles only on the first row of the grid
+            label_num = (l + 1) if one_indexed else l
+            ax.set_title(f"Layer {label_num}", fontsize=10)
+
+        # Hide any unused axes (when L is not a multiple of cols)
+        total_slots = r * cols
+        for empty_idx in range(L, total_slots):
+            rr = empty_idx // cols
+            cc = empty_idx % cols
+            axes[rr, cc].axis("off")
+
+        plt.tight_layout()
+        figs.append((fig, axes))
+    return figs

utils/modifiedGPT2.py

@@ -0,0 +1,712 @@
+from typing import Optional, Union
+import torch
+import torch.nn.functional as F
+from transformers import GPT2LMHeadModel, GPT2Model, GPT2Config
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+)
+from transformers.utils import (
+    logging,
+)
+
+from transformers.models.gpt2.modeling_gpt2 import GPT2Block, GPT2Attention, eager_attention_forward
+from torch import nn
+from typing import Callable
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+import matplotlib.pyplot as plt
+
+logger = logging.get_logger(__name__)
+
+class GPT2AttentionModified(GPT2Attention):
+    def __init__(self, config, is_cross_attention=False, layer_idx=None):
+        super().__init__(config, is_cross_attention=is_cross_attention, layer_idx=layer_idx)
+        self.config = config
+        max_positions = 2048
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
+                1, 1, max_positions, max_positions
+            ),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: Optional[tuple[torch.FloatTensor]],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ) -> tuple[Union[torch.Tensor, tuple[torch.Tensor]], ...]:
+        is_cross_attention = encoder_hidden_states is not None
+        if past_key_values is not None:
+            if isinstance(past_key_values, EncoderDecoderCache):
+                is_updated = past_key_values.is_updated.get(self.layer_idx)
+                if is_cross_attention:
+                    # after the first generated id, we can subsequently re-use all key/value_layer from cache
+                    curr_past_key_value = past_key_values.cross_attention_cache
+                else:
+                    curr_past_key_value = past_key_values.self_attention_cache
+            else:
+                curr_past_key_value = past_key_values
+
+        if is_cross_attention:
+            if not hasattr(self, "q_attn"):
+                raise ValueError(
+                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
+                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
+                )
+            query_states = self.q_attn(hidden_states)
+            attention_mask = encoder_attention_mask
+
+            # Try to get key/value states from cache if possible
+            if past_key_values is not None and is_updated:
+                key_states = curr_past_key_value.layers[self.layer_idx].keys
+                value_states = curr_past_key_value.layers[self.layer_idx].values
+            else:
+                key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
+                shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
+                key_states = key_states.view(shape_kv).transpose(1, 2)
+                value_states = value_states.view(shape_kv).transpose(1, 2)
+        else:
+            query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
+            shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
+            key_states = key_states.view(shape_kv).transpose(1, 2)
+            value_states = value_states.view(shape_kv).transpose(1, 2)
+
+        shape_q = (*query_states.shape[:-1], -1, self.head_dim)
+        query_states = query_states.view(shape_q).transpose(1, 2)
+
+        if (past_key_values is not None and not is_cross_attention) or (
+            past_key_values is not None and is_cross_attention and not is_updated
+        ):
+            # save all key/value_layer to cache to be re-used for fast auto-regressive generation
+            cache_position = cache_position if not is_cross_attention else None
+            key_states, value_states = curr_past_key_value.update(
+                key_states, value_states, self.layer_idx, {"cache_position": cache_position}
+            )
+            # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+            if is_cross_attention:
+                past_key_values.is_updated[self.layer_idx] = True
+
+        is_causal = attention_mask is None and query_states.shape[-2] > 1 and not is_cross_attention
+
+        using_eager = self.config._attn_implementation == "eager"
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        if using_eager and self.reorder_and_upcast_attn:
+            attn_output, attn_weights = self._upcast_and_reordered_attn(
+                query_states, key_states, value_states, attention_mask, head_mask
+            )
+        else:
+            if getattr(self.config, "prefix_allowed_length", None) is not None:
+                temp = self
+                temp.is_cross_attention = True
+            attn_output, attn_weights = attention_interface(
+                self if getattr(self.config, "prefix_allowed_length", None) is None else temp,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                head_mask=head_mask,
+                dropout=self.attn_dropout.p if self.training else 0.0,
+                is_causal=is_causal if getattr(self.config, "is_prefix", None) is None else False,
+                **kwargs,
+            )
+            if getattr(self.config, "plot_attention_map", False) and self.layer_idx in getattr(self.config, "plot_attention_map_layer", []):
+                # pick batch=0, head=0
+                attn_bh = attn_weights[0, 0]                 # [L,S]
+                L, S = attn_bh.shape
+                if L > 1:
+                    if getattr(self.config, "plot_attention_map_generation", 0) == 0:
+                        print(f"Plotting attention map for inputs on layer {self.layer_idx}")
+                        # full 2D heatmap
+                        data = attn_bh.detach().float().cpu().numpy()  # [L,S]
+                        plt.figure(figsize=(6,5))
+                        plt.imshow(data, aspect="auto", cmap="hot", vmin=0, vmax=0.01)
+                        plt.colorbar()
+                        plt.xlabel("Keys (S)")
+                        plt.ylabel("Queries (L)")
+                        plt.title(f"Attention map (B0,H0)  L={L}, S={S}")
+                        plt.show()
+                else:
+                    if getattr(self.config, "plot_attention_map_generation", 0) == S:
+                        print(f"Plotting attention row map for token {S} generation on layer {self.layer_idx}")
+                        # attn_bh expected shape: [..., S] for the selected (B0, H0) row
+                        row = attn_bh[0].detach().float().cpu().numpy()  # -> np.ndarray shape [S]
+                        n = row.shape[0]
+
+                        # ----- First 1024 as 32x32 -----
+                        head_1024 = row[:min(1024, n)]
+                        grid = head_1024.reshape(32, 32)
+
+                        plt.figure(figsize=(6, 5))
+                        plt.imshow(grid, aspect="auto", cmap="hot", vmin=0, vmax=0.01)
+                        plt.yticks([])
+                        plt.colorbar()
+                        plt.xlabel("Keys (S) [indices 0..1023]")
+                        plt.title(f"Attention row (B0,H0)  L={self.layer_idx}, S={S} — first 1024")
+                        plt.tight_layout()
+                        plt.show()
+
+                        # ----- Tail (>=1024) as a single-row heatmap -----
+                        tail = row[1024:]
+                        if tail.size > 0:
+                            plt.figure(figsize=(10, 1.2))
+                            # one-row heatmap
+                            plt.imshow(tail[None, :], aspect="auto", cmap="hot", vmin=0, vmax=0.01)
+                            plt.yticks([])
+                            plt.colorbar()
+                            plt.xlabel(f"Keys (S) [indices 1024..{n-1}]")
+                            plt.title(f"Attention row tail (B0,H0)  L={self.layer_idx}, S={S}")
+                            plt.tight_layout()
+                            plt.show()
+
+        attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
+        attn_output = self.c_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+
+        return attn_output, attn_weights
+
+class GPT2BlockModified(GPT2Block):
+    def __init__(self, config, layer_idx=None):
+        super().__init__(config=config)
+        self.attn = GPT2AttentionModified(config=config, layer_idx=layer_idx)
+
+    def forward(
+        self,
+        hidden_states: Optional[tuple[torch.FloatTensor]],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ) -> Union[tuple[torch.Tensor], Optional[tuple[torch.Tensor, tuple[torch.FloatTensor, ...]]]]:
+        residual = hidden_states
+        hidden_states = self.ln_1(hidden_states)
+        attn_output, self_attn_weights = self.attn(
+            hidden_states,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+
+        # residual connection
+        hidden_states = attn_output + residual
+
+        if encoder_hidden_states is not None:
+            # add one self-attention block for cross-attention
+            if not hasattr(self, "crossattention"):
+                raise ValueError(
+                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
+                    "cross-attention layers by setting `config.add_cross_attention=True`"
+                )
+            residual = hidden_states
+            hidden_states = self.ln_cross_attn(hidden_states)
+            cross_attn_output, cross_attn_weights = self.crossattention(
+                hidden_states,
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                output_attentions=output_attentions,
+            )
+            # residual connection
+            hidden_states = residual + cross_attn_output
+
+        residual = hidden_states
+        hidden_states = self.ln_2(hidden_states)
+        feed_forward_hidden_states = self.mlp(hidden_states)
+        # residual connection
+        hidden_states = residual + feed_forward_hidden_states
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+            if encoder_hidden_states is not None:
+                outputs += (cross_attn_weights,)
+
+        return outputs
+
+
+class GPT2ModelModified(GPT2Model):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.config_causal = config
+        self.config_causal._attn_implementation = "eager"  # Ensure causal mask creation uses eager implementation
+        # TEMPORARY: override the transformer blocks to pass segmentation masks
+        self.h = nn.ModuleList([GPT2BlockModified(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[tuple[tuple[torch.Tensor]], Cache]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        segmentation_mask: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:
+        r"""
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
+            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input
+            sequence tokens in the vocabulary.
+
+            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.view(-1, input_shape[-1])
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # based on pattern from src/transformers/models/whisper/modeling_whisper.py::WhisperDecoder
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            elif isinstance(past_key_values, tuple):
+                logger.warning_once(
+                    "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.53.0. "
+                    "You should pass an instance of `Cache` instead, e.g. "
+                    "`past_key_values=DynamicCache.from_legacy_cache(past_key_values)`."
+                )
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+
+            if self.config.add_cross_attention and not isinstance(past_key_values, EncoderDecoderCache):
+                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache())
+
+        if inputs_embeds is None:
+            inputs_embeds = self.wte(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        position_embeds = self.wpe(position_ids)
+        hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)
+
+        # Attention mask.
+        # ._update_causal_mask() and ._prepare_4d_causal_attention_mask_with_cache_position() copied from LlamaModel
+        if attention_mask is not None and attention_mask.ndim < 4:
+            attention_mask = attention_mask.view(batch_size, -1)
+
+        causal_mask = create_causal_mask(
+            config=self.config_causal,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
+        if self.config.add_cross_attention and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            if _use_sdpa:
+                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                    mask=encoder_attention_mask, dtype=inputs_embeds.dtype, tgt_len=input_shape[-1]
+                )
+            elif self._attn_implementation != "flash_attention_2":
+                encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # head_mask has shape n_layer x batch x n_heads x N x N
+        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+
+        if token_type_ids is not None:
+            token_type_embeds = self.wte(token_type_ids)
+            hidden_states = hidden_states + token_type_embeds
+
+        hidden_states = self.drop(hidden_states)
+
+        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
+
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+        all_hidden_states = () if output_hidden_states else None
+        for i, block in enumerate(self.h):
+            # Model parallel
+            if self.model_parallel:
+                torch.cuda.set_device(hidden_states.device)
+                if isinstance(head_mask, torch.Tensor):
+                    head_mask = head_mask.to(hidden_states.device)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            if segmentation_mask is not None and causal_mask is not None:
+                # Make a safe copy of the causal mask and ensure its spatial
+                # dimensions match the sequence length that the attention
+                # functions expect. This prevents off-by-one shape errors
+                # when using eager attention (torch.where requires same sizes).
+                causal_mask_modified = causal_mask.clone()
+                if getattr(self.config, "prefix_allowed_length", None) is not None:
+                    causal_mask_modified[:, :, :, :self.config.prefix_allowed_length].zero_()
+
+                # Use the input sequence length to crop the causal mask if needed
+                seq_len = input_shape[-1]
+                if causal_mask_modified.shape[2] != seq_len or causal_mask_modified.shape[3] != seq_len:
+                    causal_mask_modified = causal_mask_modified[:, :, :seq_len, :seq_len]
+
+                # Clip segmentation mask to fit into causal_mask_modified before adding.
+                _, _, M, N = segmentation_mask.shape
+                M = min(M, causal_mask_modified.shape[2])
+                N = min(N, causal_mask_modified.shape[3])
+                causal_mask_modified[:, :, :M, :N] += segmentation_mask[:, i, :M, :N].unsqueeze(1)
+            if getattr(self.config, "plot_attention_mask", False) and i in getattr(self.config, "plot_attention_mask_layer", [0]):
+                if segmentation_mask is not None and causal_mask is not None:
+                    print(f"Block {i}: segmentation mask added to causal mask.")
+                    plt.imshow(causal_mask_modified[0,0].detach().cpu(), aspect='auto', cmap='hot', vmin=-1, vmax=1)
+                    plt.colorbar()
+                    plt.title(f"Causal Mask with Segmentation (Block {i})")
+                    plt.show()
+                else:
+                    print(f"Block {i}: no segmentation mask applied.")
+                    plt.imshow(causal_mask[0,0].detach().cpu(), aspect='auto', cmap='hot', vmin=-1, vmax=1)
+                    plt.colorbar()
+                    plt.title(f"Causal Mask (Block {i})")
+                    plt.show()
+
+
+            outputs = block(
+                hidden_states,
+                past_key_values if not (self.gradient_checkpointing and self.training) else None,
+                cache_position,
+                causal_mask_modified if segmentation_mask is not None and causal_mask is not None else causal_mask,
+                head_mask[i],
+                encoder_hidden_states,  # as a positional argument for gradient checkpointing
+                encoder_attention_mask=encoder_attention_mask,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+
+            hidden_states = outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (outputs[1],)
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (outputs[2],)
+
+            # Model Parallel: If it's the last layer for that device, put things on the next device
+            if self.model_parallel:
+                for k, v in self.device_map.items():
+                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
+                        hidden_states = hidden_states.to("cuda:" + str(k + 1))
+
+        hidden_states = self.ln_f(hidden_states)
+
+        hidden_states = hidden_states.view(output_shape)
+        # Add last hidden state
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        past_key_values = past_key_values if use_cache else None
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions, all_cross_attentions]
+                if v is not None
+            )
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+class GPT2LMHeadModelModified(GPT2LMHeadModel):
+    def __init__(self, config):
+        super().__init__(config)
+        # replace the base transformer with our modified transformer implementation
+        self.transformer = GPT2ModelModified(config)
+        self.post_init() 
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        segmentation_mask: Optional[torch.FloatTensor] = None,
+        prefix_allowed_length: Optional[int] = None,
+        plot_attention_mask: Optional[bool] = False,
+        plot_attention_mask_layer: Optional[list[int]] = [0],
+        plot_attention_map: Optional[bool] = False,
+        plot_attention_map_layer: Optional[list[int]] = [0],
+        plot_attention_map_generation: Optional[int] = 0,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:
+        r"""
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
+            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input
+            sequence tokens in the vocabulary.
+
+            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        labels (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if prefix_allowed_length is not None:
+            self.config.prefix_allowed_length = prefix_allowed_length
+
+        if plot_attention_mask is not None:
+            self.config.plot_attention_mask = plot_attention_mask
+            if plot_attention_mask_layer is not None:
+                self.config.plot_attention_mask_layer = plot_attention_mask_layer
+
+        if plot_attention_map is not None:
+            if plot_attention_map_layer is not None:
+                self.config.plot_attention_map_layer = plot_attention_map_layer
+            if plot_attention_map_generation is not None:
+                self.config.plot_attention_map_generation = plot_attention_map_generation
+            self.config.plot_attention_map = plot_attention_map
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            segmentation_mask=segmentation_mask, #Added this parameter
+            **kwargs,
+        )
+        hidden_states = transformer_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            # Flatten the tokens
+            loss = self.loss_function(
+                logits,
+                labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+        if not return_dict:
+            output = (logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+@torch.no_grad()
+def expand_gpt2_positional_embeddings(
+    model: torch.nn.Module,
+    new_max_positions: int,
+    mode: str = "linear",        # "linear" | "copy_last" | "zeros"
+    align_corners: bool = True,  # for linear interpolation
+):
+    """
+    Expand GPT-2's learned positional embeddings (wpe) to `new_max_positions`.
+
+    Works with GPT2LMHeadModel or GPT2Model (HF). Updates model.config.n_positions (and n_ctx if present).
+    Does NOT mutate token embeddings; only position table + config.
+
+    Args:
+        model: HF GPT2LMHeadModel or GPT2Model (already loaded).
+        new_max_positions: int, desired max sequence length (e.g., 1536 or 2048).
+        mode: how to initialize new rows if expanding:
+              - "linear": 1D linear interpolation along position dim (recommended)
+              - "copy_last": copy the last learned vector into all new rows
+              - "zeros": initialize new rows to zero
+        align_corners: passed to F.interpolate for "linear" mode.
+
+    Returns:
+        model (same instance) with expanded wpe and updated config.
+    """
+    # Locate the position embedding table.
+    # Support both:
+    # - GPT2LMHeadModel (has .transformer which is a GPT2Model with .wpe)
+    # - GPT2Model (exposes .wpe directly)
+    if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
+        model_for_wpe = model.transformer
+    elif hasattr(model, "wpe"):
+        model_for_wpe = model
+    else:
+        raise ValueError("Model does not look like a GPT-2 family model with a position embedding 'wpe')")
+
+    wpe = model_for_wpe.wpe
+
+    old_n, d = wpe.weight.shape
+    if new_max_positions <= 0:
+        raise ValueError("new_max_positions must be positive")
+    if new_max_positions == old_n:
+        # Still update config for consistency
+        if hasattr(model.config, "n_positions"):
+            model.config.n_positions = new_max_positions
+        if hasattr(model.config, "n_ctx"):
+            model.config.n_ctx = new_max_positions
+        return model
+
+    device = wpe.weight.device
+    dtype  = wpe.weight.dtype
+
+    if new_max_positions < old_n:
+        # Shrink (rare): just slice
+        new_weight = wpe.weight[:new_max_positions].clone()
+    else:
+        # Expand
+        if mode == "linear":
+            # Interpolate along position dimension.
+            # Treat embedding dim as channels: (1, d, old_n) -> (1, d, new_n) -> (new_n, d)
+            w = wpe.weight.transpose(0, 1).unsqueeze(0)  # (1, d, old_n)
+            w_new = F.interpolate(w, size=new_max_positions, mode="linear", align_corners=align_corners)
+            new_weight = w_new.squeeze(0).transpose(0, 1).contiguous()  # (new_n, d)
+        elif mode == "copy_last":
+            new_weight = torch.empty((new_max_positions, d), device=device, dtype=dtype)
+            new_weight[:old_n].copy_(wpe.weight)
+            new_weight[old_n:].copy_(wpe.weight[old_n - 1].expand(new_max_positions - old_n, d))
+        elif mode == "zeros":
+            new_weight = torch.zeros((new_max_positions, d), device=device, dtype=dtype)
+            new_weight[:old_n].copy_(wpe.weight)
+        else:
+            raise ValueError(f"Unknown mode '{mode}'")
+
+    # Replace embedding module on whichever object held the original table
+    new_wpe = torch.nn.Embedding(new_max_positions, d, device=device, dtype=dtype)
+    new_wpe.weight.copy_(new_weight)
+
+    # Keep requires_grad True (default). If you want to freeze, set .requires_grad_(False).
+    if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
+        model.transformer.wpe = new_wpe
+    else:
+        model.wpe = new_wpe
+
+    # Update config fields used by HF
+    if hasattr(model.config, "n_positions"):
+        model.config.n_positions = new_max_positions
+    if hasattr(model.config, "n_ctx"):
+        model.config.n_ctx = new_max_positions
+
+    return model
+
+def create_decoder(attention = "sdpa"):
+    config = GPT2Config.from_pretrained("gpt2")
+    config._attn_implementation = attention
+    new_max_positions = 2048
+    decoder = GPT2LMHeadModelModified.from_pretrained("gpt2", config=config)
+    decoder.config._attn_implementation = attention
+    decoder = expand_gpt2_positional_embeddings(decoder, new_max_positions=new_max_positions, mode="linear")
+    return decoder

utils/processing.py

@@ -0,0 +1,27 @@
+import torchvision.transforms as T
+import fsspec
+import io
+from PIL import Image
+
+def image_transform(img_size=512):
+    return T.Compose([
+        T.Resize((img_size, img_size), interpolation=T.InterpolationMode.BICUBIC),
+        T.ToTensor(),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+def open_binary(path: str):
+    """
+    Open any (local or gs://) file for binary reading.
+    Returns a file-like object (context manager).
+    """
+    return fsspec.open(path, mode="rb").open()
+
+def pil_from_path(path: str) -> Image.Image:
+    """
+    Load an image from local or GCS; returns a PIL image in RGB.
+    """
+    with open_binary(path) as f:
+        img_bytes = f.read()
+    im = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    return im