finished

app.py

@@ -22,14 +22,15 @@ def set_seed(model_seed = 42):
     torch.cuda.manual_seed(model_seed) if torch.cuda.is_available() else None
 
 set_seed()
-clip_utils = Clip_Utils()
-clip_utils.init_Clip()
 model_utils, vl_gpt, tokenizer = None, None, None
-model_name = "Clip"
+model_utils = ChartGemma_Utils()
+vl_gpt, tokenizer = model_utils.init_ChartGemma()
+for layer in vl_gpt.language_model.model.layers:
+    layer.self_attn = ModifiedGemmaAttention(layer.self_attn)
+model_name = "ChartGemma-3B"
 language_model_max_layer = 24
-language_model_best_layer_min = 8
-language_model_best_layer_max = 8
-vision_model_best_layer = 24
+language_model_best_layer_min = 9
+language_model_best_layer_max = 15
 
 def clean():
     global model_utils, vl_gpt, tokenizer, clip_utils
@@ -71,123 +72,83 @@ def multimodal_understanding(model_type,
 
     input_text_decoded = ""
     answer = ""
-    if model_name == "Clip":
+    
         
-        inputs = clip_utils.prepare_inputs([question], image)
+    for param in vl_gpt.parameters():
+        param.requires_grad = True
 
 
-        if activation_map_method == "GradCAM":
-            # Generate Grad-CAM
-            all_layers = [layer.layer_norm1 for layer in clip_utils.model.vision_model.encoder.layers]
+    prepare_inputs = model_utils.prepare_inputs(question, image)
 
-            if visualization_layer_min != visualization_layer_max:
-                target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max-1]
-            else:
-                target_layers = [all_layers[visualization_layer_min-1]]
-            grad_cam = VisualizationClip(clip_utils.model, target_layers)
-            cam, outputs, grid_size = grad_cam.generate_cam(inputs, target_token_idx=0, visual_method=visual_method)
-            cam = cam.to("cpu")
-            cam = [generate_gradcam(cam, image, size=(224, 224))]
-            grad_cam.remove_hooks()
-            target_token_decoded = ""
-            
-    
-
-    else:
-        
-        for param in vl_gpt.parameters():
-            param.requires_grad = True
+    if response_type == "answer + visualization":
+        if model_name.split('-')[0] == "Janus":
+            inputs_embeds = model_utils.generate_inputs_embeddings(prepare_inputs)
+            outputs = model_utils.generate_outputs(inputs_embeds, prepare_inputs, temperature, top_p)
+        else:
+            outputs = model_utils.generate_outputs(prepare_inputs, temperature, top_p)
 
+        sequences = outputs.sequences.cpu().tolist()
+        answer = tokenizer.decode(sequences[0], skip_special_tokens=True)
+        attention_raw = outputs.attentions
+        print("answer generated")
 
-        prepare_inputs = model_utils.prepare_inputs(question, image)
+    input_ids = prepare_inputs.input_ids[0].cpu().tolist()
+    input_ids_decoded = [tokenizer.decode([input_ids[i]]) for i in range(len(input_ids))]
 
-        if response_type == "answer + visualization":
-            if model_name.split('-')[0] == "Janus":
-                inputs_embeds = model_utils.generate_inputs_embeddings(prepare_inputs)
-                outputs = model_utils.generate_outputs(inputs_embeds, prepare_inputs, temperature, top_p)
-            else:
-                outputs = model_utils.generate_outputs(prepare_inputs, temperature, top_p)
-
-            sequences = outputs.sequences.cpu().tolist()
-            answer = tokenizer.decode(sequences[0], skip_special_tokens=True)
-            attention_raw = outputs.attentions
-            print("answer generated")
-
-        input_ids = prepare_inputs.input_ids[0].cpu().tolist()
-        input_ids_decoded = [tokenizer.decode([input_ids[i]]) for i in range(len(input_ids))]
-
-        if activation_map_method == "GradCAM":
-            # target_layers = vl_gpt.vision_model.vision_tower.blocks
-            if focus == "Visual Encoder":
-                if model_name.split('-')[0] == "Janus":
-                    all_layers = [block.norm1 for block in vl_gpt.vision_model.vision_tower.blocks]
-                else:
-                    all_layers = [block.layer_norm1 for block in vl_gpt.vision_tower.vision_model.encoder.layers]
-            else:
-                all_layers = [layer.self_attn for layer in vl_gpt.language_model.model.layers]
-            
-            print("layer values:", visualization_layer_min, visualization_layer_max)
-            if visualization_layer_min != visualization_layer_max:
-                print("multi layers")
-                target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max]
-            else:
-                print("single layer")
-                target_layers = [all_layers[visualization_layer_min-1]]
-            
+    if activation_map_method == "AG-CAM":
+        # target_layers = vl_gpt.vision_model.vision_tower.blocks
+        
+        all_layers = [layer.self_attn for layer in vl_gpt.language_model.model.layers]
+        
+        print("layer values:", visualization_layer_min, visualization_layer_max)
+        if visualization_layer_min != visualization_layer_max:
+            print("multi layers")
+            target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max]
+        else:
+            print("single layer")
+            target_layers = [all_layers[visualization_layer_min-1]]
+        
 
-            if model_name.split('-')[0] == "Janus":
-                gradcam = VisualizationJanus(vl_gpt, target_layers)
-            elif model_name.split('-')[0] == "LLaVA":
-                gradcam = VisualizationLLaVA(vl_gpt, target_layers)
-            elif model_name.split('-')[0] == "ChartGemma":
-                gradcam = VisualizationChartGemma(vl_gpt, target_layers)
+        if model_name.split('-')[0] == "Janus":
+            gradcam = VisualizationJanus(vl_gpt, target_layers)
+        elif model_name.split('-')[0] == "LLaVA":
+            gradcam = VisualizationLLaVA(vl_gpt, target_layers)
+        elif model_name.split('-')[0] == "ChartGemma":
+            gradcam = VisualizationChartGemma(vl_gpt, target_layers)
 
-            start = 0
-            cam = []
-            if focus == "Visual Encoder":
-                if target_token_idx != -1:
-                    cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, focus)
-                    cam_grid = cam_tensors.reshape(grid_size, grid_size)
-                    cam_i = generate_gradcam(cam_grid, image)
-                    cam_i = add_title_to_image(cam_i, input_ids_decoded[start + target_token_idx])
-                    cam = [cam_i]
-                else:
-                    i = 0
-                    cam = []
-                    while start + i < len(input_ids_decoded):
-                        if model_name.split('-')[0] == "Janus":
-                            gradcam = VisualizationJanus(vl_gpt, target_layers)
-                        elif model_name.split('-')[0] == "LLaVA":
-                            gradcam = VisualizationLLaVA(vl_gpt, target_layers)
-                        elif model_name.split('-')[0] == "ChartGemma":
-                            gradcam = VisualizationChartGemma(vl_gpt, target_layers)
-                        cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_method, focus, accumulate_method)
-                        cam_grid = cam_tensors.reshape(grid_size, grid_size)
-                        cam_i = generate_gradcam(cam_grid, image)
-                        cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
-                        cam.append(cam_i)
-                        gradcam.remove_hooks()
-                        i += 1
+        start = 0
+        cam = []
+        
+        # utilize the entire sequence, including <image>s, question, and answer
+        entire_inputs = prepare_inputs
+        if response_type == "answer + visualization" and focus == "question + answer":
+            if model_name.split('-')[0] == "Janus" or model_name.split('-')[0] == "LLaVA":
+                entire_inputs = model_utils.prepare_inputs(question, image, answer)
             else:
-                cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, focus, accumulate_method)
-                if target_token_idx != -1:
-                    input_text_decoded = input_ids_decoded[start + target_token_idx]
-                    for i, cam_tensor in enumerate(cam_tensors):
-                        if i == target_token_idx:
-                            cam_grid = cam_tensor.reshape(grid_size, grid_size)
-                            cam_i = generate_gradcam(cam_grid, image)
-                            cam = [add_title_to_image(cam_i, input_text_decoded)]
-                            break
-                else:
-                    cam = []
-                    for i, cam_tensor in enumerate(cam_tensors):
-                        cam_grid = cam_tensor.reshape(grid_size, grid_size)
-                        cam_i = generate_gradcam(cam_grid, image)
-                        cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
-
-                        cam.append(cam_i)
-                     
-            gradcam.remove_hooks()
+                entire_inputs["input_ids"] = outputs.sequences
+                entire_inputs["attention_mask"] = torch.ones_like(outputs.sequences)
+            input_ids = entire_inputs['input_ids'][0].cpu().tolist()
+            input_ids_decoded = [tokenizer.decode([input_ids[i]]) for i in range(len(input_ids))]
+
+        cam_tensors, grid_size, start = gradcam.generate_cam(entire_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, "Language Model", accumulate_method)
+        if target_token_idx != -1:
+            input_text_decoded = input_ids_decoded[start + target_token_idx]
+            for i, cam_tensor in enumerate(cam_tensors):
+                if i == target_token_idx:
+                    cam_grid = cam_tensor.reshape(grid_size, grid_size)
+                    cam_i = generate_gradcam(cam_grid, image)
+                    cam = [add_title_to_image(cam_i, input_text_decoded)]
+                    break
+        else:
+            cam = []
+            for i, cam_tensor in enumerate(cam_tensors):
+                cam_grid = cam_tensor.reshape(grid_size, grid_size)
+                cam_i = generate_gradcam(cam_grid, image)
+                cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
+
+                cam.append(cam_i)
+                    
+        gradcam.remove_hooks()
                 
 
     # Collect Results
@@ -219,34 +180,7 @@ def model_slider_change(model_type):
     global model_utils, vl_gpt, tokenizer, clip_utils, model_name, language_model_max_layer, language_model_best_layer_min, language_model_best_layer_max, vision_model_best_layer
     model_name = model_type
 
-
-    encoder_only_res = [
-        gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type"),
-        gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus"),
-        gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-        gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
-    ]
-
-    language_res = [
-        gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
-        gr.Dropdown(choices=["Language Model"], value="Language Model", label="focus"),
-        gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-        gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
-    ]
-
-
-    if model_type == "Clip":
-        clean()
-        set_seed()
-        clip_utils = Clip_Utils()
-        clip_utils.init_Clip()
-        sliders = [
-            gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min"),
-            gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max"),
-        ]
-        return tuple(encoder_only_res + sliders)
-    
-    elif model_type.split('-')[0] == "Janus":
+    if model_type.split('-')[0] == "Janus":
         # best seed: 70
         clean()
         set_seed()
@@ -263,7 +197,7 @@ def model_slider_change(model_type):
             gr.Slider(minimum=1, maximum=24, value=language_model_best_layer_min, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=24, value=language_model_best_layer_max, step=1, label="visualization layers max"),
         ]
-        return tuple(language_res + sliders)
+        return tuple(sliders)
     
     elif model_type.split('-')[0] == "LLaVA":
         
@@ -280,7 +214,7 @@ def model_slider_change(model_type):
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_min, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_max, step=1, label="visualization layers max"),
         ]
-        return tuple(language_res + sliders)
+        return tuple(sliders)
     
     elif model_type.split('-')[0] == "ChartGemma":
         clean()
@@ -290,62 +224,16 @@ def model_slider_change(model_type):
         for layer in vl_gpt.language_model.model.layers:
             layer.self_attn = ModifiedGemmaAttention(layer.self_attn)
         language_model_max_layer = 18
-        vision_model_best_layer = 19
-        language_model_best_layer_min = 11
+        language_model_best_layer_min = 9
         language_model_best_layer_max = 15
 
         sliders = [
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_min, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_max, step=1, label="visualization layers max"),
         ]
-        return tuple(language_res + sliders)
-
-    
+        return tuple(sliders)
 
 
-def focus_change(focus):
-    global model_name, language_model_max_layer
-    if model_name == "Clip":
-        res = (
-                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-                gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max")
-            )
-        return res
-
-    if focus == "Language Model":
-        if response_type.value == "answer + visualization":
-            res = (
-                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-                gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_min, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_max, step=1, label="visualization layers max")
-            )
-            return res
-        else:
-            res = (
-                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-                gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_min, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer_max, step=1, label="visualization layers max")
-            )
-            return res
-
-    else:
-        if model_name.split('-')[0] == "ChartGemma":
-            res = (
-                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers max")
-            )
-            return res
-        
-        else:
-            res = (
-                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
-                gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers max")
-            )
-            return res
-
 
 def test_change(test_selector):
     if test_selector == "mini-VLAT":
@@ -376,7 +264,7 @@ with gr.Blocks() as demo:
     with gr.Row():
 
         with gr.Column():
-            model_selector = gr.Dropdown(choices=["Clip", "ChartGemma-3B", "Janus-Pro-1B", "Janus-Pro-7B", "LLaVA-1.5-7B"], value="Clip", label="model")
+            model_selector = gr.Dropdown(choices=["ChartGemma-3B", "Janus-Pro-1B", "Janus-Pro-7B", "LLaVA-1.5-7B"], value="ChartGemma-3B", label="model")
             test_selector = gr.Dropdown(choices=["mini-VLAT", "VLAT", "VLAT-old"], value="mini-VLAT", label="test")
             question_input = gr.Textbox(label="Input Prompt")
             und_seed_input = gr.Number(label="Seed", precision=0, value=42)
@@ -386,15 +274,15 @@ with gr.Blocks() as demo:
 
         
         with gr.Column():
-            response_type = gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type")
-            focus = gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus")
-            activation_map_method = gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="visualization type")
+            response_type = gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type")
+            focus = gr.Dropdown(choices=["question", "question + answer"], value="question + answer", label="focus")
+            activation_map_method = gr.Dropdown(choices=["AG-CAM"], value="AG-CAM", label="visualization type")
             accumulate_method = gr.Dropdown(choices=["sum", "mult"], value="sum", label="layers accumulate method")
-            visual_method = gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
+            visual_method = gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
             
 
-            visualization_layers_min = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min")
-            visualization_layers_max = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max")
+            visualization_layers_min = gr.Slider(minimum=1, maximum=18, value=11, step=1, label="visualization layers min")
+            visualization_layers_max = gr.Slider(minimum=1, maximum=18, value=15, step=1, label="visualization layers max")
 
         
         
@@ -404,24 +292,10 @@ with gr.Blocks() as demo:
             fn=model_slider_change, 
             inputs=model_selector, 
             outputs=[
-                response_type,
-                focus,
-                activation_map_method,
-                visual_method,
                 visualization_layers_min,
                 visualization_layers_max
             ]
         )
-        
-        focus.change(
-            fn = focus_change,
-            inputs = focus,
-            outputs=[
-                activation_map_method,
-                visualization_layers_min,
-                visualization_layers_max,
-            ]
-        )
 
         
 

demo/model_utils.py

@@ -74,14 +74,14 @@ class Janus_Utils(Model_Utils):
         return self.vl_gpt, self.tokenizer
     
     @spaces.GPU(duration=120)
-    def prepare_inputs(self, question, image):
+    def prepare_inputs(self, question, image, answer=None):
         conversation = [
             {
                 "role": "<|User|>",
                 "content": f"<image_placeholder>\n{question}",
                 "images": [image],
             },
-            {"role": "<|Assistant|>", "content": ""},
+            {"role": "<|Assistant|>", "content": answer if answer else ""}
         ]
         
         pil_images = [Image.fromarray(image)]
@@ -152,16 +152,33 @@ class LLaVA_Utils(Model_Utils):
         return self.vl_gpt, self.tokenizer
     
     @spaces.GPU(duration=120)
-    def prepare_inputs(self, question, image):
-        conversation = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": question},
-                    {"type": "image"},
-                ],
-            },
-        ]
+    def prepare_inputs(self, question, image, answer=None):
+        if answer:
+            conversation = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": question},
+                        {"type": "image"},
+                    ],
+                },
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "text", "text": answer},
+                    ],
+                }
+            ]
+        else:
+            conversation = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": question},
+                        {"type": "image"},
+                    ],
+                },
+            ]
         
         prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
         pil_images = [Image.fromarray(image).resize((384, 384))]

demo/visualization.py

@@ -25,7 +25,7 @@ class Visualization:
             self.hooks.append(layer.register_backward_hook(self._backward_hook))
 
     def _forward_hook(self, module, input, output):
-        print("forward_hook: self_attn_input: ", input)
+        # print("forward_hook: self_attn_input: ", input)
         self.activations.append(output)
 
     def _backward_hook(self, module, grad_in, grad_out):
@@ -42,12 +42,12 @@ class Visualization:
             layer.get_attn_map = types.MethodType(get_attn_map, layer)
 
     def _forward_activate_hooks(self, module, input, output):
-        print("forward_activate_hool: module: ", module)
-        print("forward_activate_hook: self_attn_input: ", input)
+        # print("forward_activate_hool: module: ", module)
+        # print("forward_activate_hook: self_attn_input: ", input)
 
         attn_output, attn_weights = output  # Unpack outputs
-        print("attn_output shape:", attn_output.shape)
-        print("attn_weights shape:", attn_weights.shape)
+        # print("attn_output shape:", attn_output.shape)
+        # print("attn_weights shape:", attn_weights.shape)
         module.save_attn_map(attn_weights)
         attn_weights.register_hook(module.save_attn_gradients)
 
@@ -137,8 +137,10 @@ class Visualization:
             
             grad = F.relu(grad)
 
+            # cam = act
             # cam = grad
             cam = act * grad # shape: [1, heads, seq_len, seq_len]
+
             cam = cam.sum(dim=1) # shape: [1, seq_len, seq_len]
             cam = cam.to(torch.float32).detach().cpu()
             cams.append(cam)
@@ -187,7 +189,6 @@ class Visualization:
             # print("cam_sum shape: ", cam_sum.shape)
             num_patches = cam_sum.shape[-1]  # Last dimension of CAM output
             grid_size = int(num_patches ** 0.5)
-            # print(f"Detected grid size: {grid_size}x{grid_size}")
 
             cam_sum = cam_sum.view(grid_size, grid_size)
             if normalize:
@@ -207,7 +208,6 @@ class Visualization:
 
                 num_patches = cam_l_i.shape[-1]  # Last dimension of CAM output
                 grid_size = int(num_patches ** 0.5)
-                # print(f"Detected grid size: {grid_size}x{grid_size}")
 
                 # Fix the reshaping step dynamically
                 cam_reshaped = cam_l_i.view(grid_size, grid_size)
@@ -258,7 +258,6 @@ class VisualizationClip(Visualization):
         
     @spaces.GPU(duration=120)
     def generate_cam(self, input_tensor, target_token_idx=None, visual_method="CLS"):
-        """ Generates Grad-CAM heatmap for ViT. """
         self.setup_grads()
         # Forward Backward pass
         output_full = self.forward_backward(input_tensor, visual_method, target_token_idx)
@@ -301,9 +300,13 @@ class VisualizationJanus(Visualization):
     def forward_backward(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder"):
         # Forward
         image_embeddings, inputs_embeddings, outputs = self.model(input_tensor, tokenizer, temperature, top_p)
-        input_ids = input_tensor.input_ids
+        print(input_tensor.keys())
+        input_ids = input_tensor["input_ids"]
         start_idx = 620
         self.model.zero_grad()
+
+        
+        
         if focus == "Visual Encoder":
             loss = outputs.logits.max(dim=-1).values[0, start_idx + target_token_idx]
             loss.backward()
@@ -335,11 +338,15 @@ class VisualizationJanus(Visualization):
 
         elif focus == "Language Model":
             
-            cam_sum = self.grad_cam_llm(mean_inside=True)
+            # cam_sum = self.grad_cam_llm(mean_inside=True)
             
-            images_seq_mask = input_tensor.images_seq_mask
+            images_seq_mask = input_tensor.images_seq_mask[0].detach().cpu().tolist()
             
-            cam_sum_lst, grid_size = self.process_multiple(cam_sum, start_idx, images_seq_mask)
+            # cam_sum_lst, grid_size = self.process_multiple(cam_sum, start_idx, images_seq_mask)
+
+            cams = self.attn_guided_cam()
+            cam_sum_lst, grid_size = self.process_multiple_acc(cams, start_idx, images_seq_mask, accumulate_method=accumulate_method)
+
 
             return cam_sum_lst, grid_size, start_idx
 
@@ -407,15 +414,6 @@ class VisualizationChartGemma(Visualization):
         self._modify_layers()
         self._register_hooks_activations()
 
-    # def custom_loss(self, start_idx, input_ids, logits):
-    #     Q = logits.shape[1]
-    #     loss = 0
-    #     q = 0
-    #     while start_idx + q < Q - 1:
-    #         loss += F.cross_entropy(logits[0, start_idx + q], input_ids[0, start_idx + q + 1])
-    #         q += 1
-    #     return loss
-
     
     def forward_backward(self, inputs, focus, start_idx, target_token_idx, visual_method="softmax"):
         outputs_raw = self.model(**inputs, output_hidden_states=True)
@@ -429,9 +427,11 @@ class VisualizationChartGemma(Visualization):
         elif focus == "Language Model":
             self.model.zero_grad()
             print("logits shape:", outputs_raw.logits.shape)
+            print("start_idx:", start_idx)
             if target_token_idx == -1:
-                loss = outputs_raw.logits.max(dim=-1).values.sum()
-                # loss = self.custom_loss(start_idx, inputs['input_ids'], outputs_raw.logits)
+                logits_prob = F.softmax(outputs_raw.logits, dim=-1)
+                loss = logits_prob.max(dim=-1).values.sum()
+
             else:
                 loss = outputs_raw.logits.max(dim=-1).values[0, start_idx + target_token_idx]
             loss.backward()
@@ -495,7 +495,7 @@ def generate_gradcam(
     normalize=False
 ):
     """
-    Generates a Grad-CAM heatmap overlay on top of the input image.
+    Generates a heatmap overlay on top of the input image.
 
     Parameters:
         cam (torch.Tensor): A tensor of shape (C, H, W) representing the
@@ -508,9 +508,8 @@ def generate_gradcam(
         normalize (bool): Whether to normalize the heatmap (default False).
 
     Returns:
-        PIL.Image: The image overlaid with the Grad-CAM heatmap.
+        PIL.Image: The image overlaid with the heatmap.
     """
-    # print("Generating Grad-CAM with shape:", cam.shape)
 
     if normalize:
         cam_min, cam_max = cam.min(), cam.max()