models/task_feature_matrix.py

"""Task-Feature Matrix computation for AlignX Stage 1.

T_a = W1 @ compress(Δθ_a) + W2 @ F_a

Δθ_a: task vector (compressed LoRA delta to latent_dim)
F_a:  alignment feature vector (mean hidden state, size hidden_dim)
T_a:  k-dim latent alignment representation
"""

import torch
import torch.nn as nn
from typing import Dict, Optional


def compress_task_vector(
    delta_state_dict: Dict[str, torch.Tensor],
    target_dim: int = 4096,
) -> torch.Tensor:
    """Flatten and mean-pool all parameter deltas to a fixed-size vector.

    LoRA deltas are small (~4M params for rank-8) so we can flatten them
    and chunk-average down to `target_dim` to make W1 tractable.
    """
    all_params = []
    for v in delta_state_dict.values():
        if v.dtype in (torch.float32, torch.float16, torch.bfloat16):
            all_params.append(v.detach().cpu().float().flatten())

    if not all_params:
        return torch.zeros(target_dim)

    flat = torch.cat(all_params)
    total = flat.shape[0]

    if total >= target_dim:
        chunk = total // target_dim
        remainder = total - chunk * target_dim
        trimmed = flat[:chunk * target_dim]
        compressed = trimmed.view(target_dim, chunk).mean(dim=1)
    else:
        # Pad with zeros
        compressed = torch.zeros(target_dim)
        compressed[:total] = flat

    return compressed.float()


def extract_alignment_features(
    model,
    dataset,
    layer_idx: int = -1,
    device: str = "cuda",
    max_samples: int = 200,
    batch_size: int = 4,
) -> torch.Tensor:
    """Extract mean hidden-state vector F_a at transformer layer `layer_idx`.

    Runs a forward pass over `dataset` samples, collecting the hidden state
    from the chosen layer, then returns the mean across all samples and tokens.
    """
    import torch
    from torch.utils.data import DataLoader

    model.eval()
    # Only call .to(device) if the model isn't already dispatched across multiple devices
    # (dispatch_model sets model.hf_device_map; calling .to() on a dispatched model
    # would un-dispatch it and may OOM a single GPU).
    if not getattr(model, "hf_device_map", None):
        model.to(device)

    hidden_dim = model.config.hidden_size
    feature_accum = torch.zeros(hidden_dim, device=device)
    count = 0

    # Identify the target layer
    layers = _get_transformer_layers(model)
    target_layer = layers[layer_idx]

    captured = []

    def _hook(module, input, output):
        # output can be tuple (hidden_state, ...) or just tensor
        hs = output[0] if isinstance(output, tuple) else output
        # hs: (batch, seq_len, hidden_dim)
        captured.append(hs.detach().mean(dim=(0, 1)).cpu())  # (hidden_dim,)

    handle = target_layer.register_forward_hook(_hook)

    # When the model is dispatched across GPUs (hf_device_map set), inputs must go
    # to cuda:0 (where the embedding layer lives), not to the `device` argument.
    input_device = "cuda:0" if getattr(model, "hf_device_map", None) else device

    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    with torch.no_grad():
        for i, batch in enumerate(loader):
            if count >= max_samples:
                break
            input_ids = batch["input_ids"].to(input_device)
            attention_mask = batch["attention_mask"].to(input_device)
            try:
                model(input_ids=input_ids, attention_mask=attention_mask)
            except Exception:
                pass
            count += input_ids.shape[0]

    handle.remove()

    if captured:
        F_a = torch.stack(captured).mean(0)  # (hidden_dim,)
    else:
        F_a = torch.zeros(hidden_dim)

    print(f"[FeatureExtraction] Extracted F_a of shape {F_a.shape} from {count} samples")
    return F_a.float()


def _get_transformer_layers(model):
    """Return the list of transformer layers for common model families."""
    # Handles LLaMA, Mistral, Gemma, DeepSeek
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    raise ValueError(f"Cannot locate transformer layers for {type(model)}")


class TaskFeatureMatrix(nn.Module):
    """T_a = W1 @ Δθ_a + W2 @ F_a

    Maps the compressed task vector and alignment feature vector
    into a shared k-dimensional latent alignment space.
    """

    def __init__(self, k: int = 256, compress_dim: int = 4096, hidden_dim: int = 4096):
        super().__init__()
        self.k = k
        self.compress_dim = compress_dim
        self.hidden_dim = hidden_dim
        # W1: (k, compress_dim)  maps compressed delta-theta to latent space
        self.W1 = nn.Linear(compress_dim, k, bias=False)
        # W2: (k, hidden_dim)    maps feature vector to latent space
        self.W2 = nn.Linear(hidden_dim, k, bias=False)
        nn.init.xavier_uniform_(self.W1.weight)
        nn.init.xavier_uniform_(self.W2.weight)

    def forward(
        self,
        delta_theta_compressed: torch.Tensor,
        F_a: torch.Tensor,
    ) -> torch.Tensor:
        """
        delta_theta_compressed: (compress_dim,) — compressed task vector
        F_a: (hidden_dim,) — alignment feature vector
        Returns T_a: (k,)
        """
        T_a = self.W1(delta_theta_compressed) + self.W2(F_a)
        return T_a

    def compute_and_cache(
        self,
        delta_theta_compressed: torch.Tensor,
        F_a: torch.Tensor,
    ) -> torch.Tensor:
        with torch.no_grad():
            return self.forward(delta_theta_compressed, F_a)