Update app.py

app.py

@@ -165,23 +165,27 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
                 dataframes.append({"Error": [f"Layer {layer_name} not found."]})
                 html_plots.append(f"<p>Error: Layer {layer_name} not found.</p>")
 
-        # Attribution: Integrated Gradients
-        def forward_func(inputs, attention_mask=None):
-            inputs = inputs.to(dtype=torch.long)  # Ensure LongTensor
-            outputs = model(inputs, attention_mask=attention_mask)
+        # Attribution: Integrated Gradients on embeddings
+        def get_embeddings(inputs, attention_mask=None):
+            with torch.no_grad():
+                embeddings = model.bert.embeddings(inputs)  # Get float embeddings
+            return embeddings
+
+        def forward_func(embeddings, attention_mask=None):
+            outputs = model(inputs_embeds=embeddings, attention_mask=attention_mask)
             return outputs.pooler_output[:, int(attribution_target)]
 
         ig = IntegratedGradients(forward_func)
         try:
-            # Ensure input_ids is LongTensor and requires grad
-            input_ids = input_ids.to(dtype=torch.long).requires_grad_(True)
+            # Get embeddings for input_ids
+            embeddings = get_embeddings(input_ids, attention_mask).requires_grad_(True)
             attributions, _ = ig.attribute(
-                inputs=input_ids,
+                inputs=embeddings,
                 additional_forward_args=(attention_mask,),
                 target=int(attribution_target),
                 return_convergence_delta=True
             )
-            attr = attributions[0].detach().numpy().sum(axis=1)
+            attr = attributions[0].detach().numpy().sum(axis=1)  # Sum over hidden size
             attr_df = pd.DataFrame({"Token": tokens, "Attribution": attr})
             attr_df.index = [f"idx_{i}" for i in range(len(attr_df))]  # String indices
             dataframes.append(convert_dict_keys_to_str(attr_df.to_dict()))
@@ -260,7 +264,7 @@ def create_gradio_interface():
 if __name__ == "__main__":
     try:
         demo = create_gradio_interface()
-        demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+        demo.launch(server_name="0.0.0.0", server_port=7861, share=False)
     except Exception as e:
         logger.error(f"Failed to launch Gradio demo: {e}")
         print(f"Error launching demo: {e}. Try running locally with a different port or without share=True.")