Aluode/PerceptionLabPortable

1# Copyright 2025 The HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15 16import requests17from PIL import Image18 19from ..masking_utils import create_causal_mask20from ..models.auto.auto_factory import _get_model_class21from ..models.auto.configuration_auto import AutoConfig22from ..models.auto.modeling_auto import MODEL_FOR_PRETRAINING_MAPPING, MODEL_MAPPING23from ..models.auto.processing_auto import PROCESSOR_MAPPING_NAMES, AutoProcessor24from ..models.auto.tokenization_auto import TOKENIZER_MAPPING_NAMES, AutoTokenizer25from .import_utils import is_torch_available26 27 28if is_torch_available():29    import torch30    import torch.nn as nn31 32# Print the matrix with words as row labels33GREEN = "\033[92m"34YELLOW = "\033[93m"35RESET = "\033[0m"36BLACK_SQUARE = "■"37WHITE_SQUARE = "⬚"38 39 40def generate_attention_matrix_from_mask(41    words, mask, img_token="<img>", sliding_window=None, token_type_ids=None, image_seq_length=None42):43    """44    Generates an attention matrix from a given attention mask.45 46    Optionally applies a sliding window mask (e.g., for Gemma2/3) and47    marks regions where image tokens occur based on the specified `img_token`.48    """49    mask = mask.int()50    if mask.ndim == 3:51        mask = mask[0, :, :]52    if mask.ndim == 4:53        mask = mask[0, 0, :, :]54 55    n = len(words)56    max_word_length = max(len(repr(word)) for word in words)57    first_img_idx = 058    output = []59 60    for i, k in enumerate(words):61        if k == img_token and not first_img_idx:62            first_img_idx = i63            mask[i, i] = 2  # Mark yellow regions64        if first_img_idx > 0 and (k != img_token or i == n - 1):65            if i == n - 1:66                i += 167            mask[first_img_idx:i, first_img_idx:i] = 2  # Mark yellow regions68            first_img_idx = 069 70    # Generate sliding window mask (size = 4), excluding img_token71    sliding_window_mask = None72    if sliding_window is not None:73        sliding_window_mask = [[1 if (0 <= i - j < sliding_window) else 0 for j in range(n)] for i in range(n)]74 75    row_dummy = " ".join(76        f"{YELLOW}{BLACK_SQUARE}{RESET}"77        if mask[0, j]78        else f"{GREEN}{BLACK_SQUARE}{RESET}"79        if 0 == j80        else BLACK_SQUARE81        if mask[0, j]82        else WHITE_SQUARE83        for j in range(n)84    )85 86    if token_type_ids is not None:87        is_special = token_type_ids == 188        token_type_buckets = torch.where(89            (token_type_ids.cumsum(-1) % 5 + is_special).bool(), token_type_ids.cumsum(-1), 090        )91        boundaries = torch.arange(0, image_seq_length + 1, image_seq_length)92        token_type_buckets = torch.bucketize(token_type_buckets, boundaries=boundaries)93 94    # Print headers95    legend = f"{GREEN}{BLACK_SQUARE}{RESET}: i == j (diagonal)   {YELLOW}{BLACK_SQUARE}{RESET}: token_type_ids"96    output.append(" " + legend)97    f_string = " " * (max_word_length + 5) + "Attention Matrix".ljust(len(row_dummy) // 2)98    if sliding_window is not None:99        f_string += "Sliding Window Mask"100    output.append(f_string)101 102    vertical_header = []103    for idx, word in enumerate(words):104        if mask[idx, idx] == 2:105            vertical_header.append([f"{YELLOW}{k}{RESET}" for k in list(str(idx).rjust(len(str(n))))])106        else:107            vertical_header.append(list(str(idx).rjust(len(str(n)))))108 109    vertical_header = list(map(list, zip(*vertical_header)))  # Transpose110 111    for row in vertical_header:112        output.append(113            (max_word_length + 5) * " " + " ".join(row) + "    |    " + " ".join(row)114            if sliding_window is not None115            else ""116        )117    for i, word in enumerate(words):118        word_repr = repr(word).ljust(max_word_length)119        colored_word = f"{YELLOW}{word_repr}{RESET}" if img_token in word else word_repr120        row_display = " ".join(121            f"{YELLOW}{BLACK_SQUARE}{RESET}"122            if img_token in words[j] and mask[i, j] and img_token in word123            else f"{GREEN}{BLACK_SQUARE}{RESET}"124            if i == j125            else BLACK_SQUARE126            if mask[i, j]127            else WHITE_SQUARE128            for j in range(n)129        )130        sliding_window_row = ""131        if sliding_window is not None:132            sliding_window_row = " ".join(133                f"{YELLOW}{BLACK_SQUARE}{RESET}"134                if img_token in words[j] and img_token in word and token_type_buckets[0, i] == token_type_buckets[0, j]135                else f"{GREEN}{BLACK_SQUARE}{RESET}"136                if i == j137                else BLACK_SQUARE138                if sliding_window_mask[i][j]139                else WHITE_SQUARE140                for j in range(n)141            )142 143        output.append(f"{colored_word}: {str(i).rjust(2)} {row_display}    |    {sliding_window_row}")144 145    return "\n".join(output)146 147 148class AttentionMaskVisualizer:149    def __init__(self, model_name: str):150        config = AutoConfig.from_pretrained(model_name)151        self.image_token = "<img>"152        if hasattr(config.get_text_config(), "sliding_window"):153            self.sliding_window = getattr(config.get_text_config(), "sliding_window", None)154        try:155            mapped_cls = _get_model_class(config, MODEL_MAPPING)156        except Exception:157            mapped_cls = _get_model_class(config, MODEL_FOR_PRETRAINING_MAPPING)158 159        if mapped_cls is None:160            raise ValueError(f"Model name {model_name} is not supported for attention visualization")161        self.mapped_cls = mapped_cls162 163        class _ModelWrapper(mapped_cls, nn.Module):164            def __init__(self, config, model_name):165                nn.Module.__init__(self)166                self.dummy_module = nn.Linear(1, 1)167                self.config = config168 169        self.model = _ModelWrapper(config, model_name)170        self.model.to(config.dtype)171        self.repo_id = model_name172        self.config = config173 174    def __call__(self, input_sentence: str, suffix=""):175        self.visualize_attention_mask(input_sentence, suffix=suffix)176 177    def visualize_attention_mask(self, input_sentence: str, suffix=""):178        model = self.model179        kwargs = {}180        image_seq_length = None181        if self.config.model_type in PROCESSOR_MAPPING_NAMES:182            img = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg?download=true"183            img = Image.open(requests.get(img, stream=True).raw)184            image_seq_length = 5185            processor = AutoProcessor.from_pretrained(self.repo_id, image_seq_length=image_seq_length)186            if hasattr(processor, "image_token"):187                image_token = processor.image_token188            else:189                image_token = processor.tokenizer.convert_ids_to_tokens([processor.image_token_id])[0]190 191            if image_token:192                input_sentence = input_sentence.replace("<img>", image_token)193 194            inputs = processor(images=img, text=input_sentence, suffix=suffix, return_tensors="pt")195 196            self.image_token = processor.tokenizer.convert_ids_to_tokens([processor.image_token_id])[0]197 198            attention_mask = inputs["attention_mask"]199            if "token_type_ids" in inputs:  # TODO inspect signature of update causal mask200                kwargs["token_type_ids"] = inputs["token_type_ids"]201            tokens = processor.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])202        elif self.config.model_type in TOKENIZER_MAPPING_NAMES:203            tokenizer = AutoTokenizer.from_pretrained(self.repo_id)204            tokens = tokenizer.tokenize(input_sentence)205            attention_mask = tokenizer(input_sentence, return_tensors="pt")["attention_mask"]206        else:207            raise ValueError(f"Model type {model.config.model_type} does not support attention visualization")208 209        model.config._attn_implementation = "eager"210        model.train()211 212        batch_size, seq_length = attention_mask.shape213        input_embeds = torch.zeros((batch_size, seq_length, model.config.hidden_size), dtype=self.model.dtype)214        cache_position = torch.arange(seq_length)215 216        causal_mask = create_causal_mask(217            config=model.config,218            input_embeds=input_embeds,219            attention_mask=attention_mask,220            cache_position=cache_position,221            past_key_values=None,222        )223 224        if causal_mask is not None:225            attention_mask = ~causal_mask.bool()226        else:227            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).expand(batch_size, 1, seq_length, seq_length)228        top_bottom_border = "##" * (229            len(f"Attention visualization for {self.config.model_type} | {self.mapped_cls}") + 4230        )  # Box width adjusted to text length231        side_border = "##"232        print(f"\n{top_bottom_border}")233        print(234            "##"235            + f"  Attention visualization for \033[1m{self.config.model_type}:{self.repo_id}\033[0m {self.mapped_cls.__name__}".center(236                len(top_bottom_border)237            )238            + "    "239            + side_border,240        )241        print(f"{top_bottom_border}")242        f_string = generate_attention_matrix_from_mask(243            tokens,244            attention_mask,245            img_token=self.image_token,246            sliding_window=getattr(self.config, "sliding_window", None),247            token_type_ids=kwargs.get("token_type_ids"),248            image_seq_length=image_seq_length,249        )250        print(f_string)251        print(f"{top_bottom_border}")252