Update app.py

app.py

@@ -15,10 +15,10 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize BERT model and tokenizer
+# Initialize BERT model and tokenizer with eager attention
 try:
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-    model = BertModel.from_pretrained('bert-base-uncased')
+    model = BertModel.from_pretrained('bert-base-uncased', attn_implementation="eager")
     model.eval()
 except Exception as e:
     logger.error(f"Failed to load BERT model: {e}")
@@ -67,7 +67,7 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
 
         # Tokenize input
         inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=512)
-        input_ids = inputs['input_ids']
+        input_ids = inputs['input_ids'].to(dtype=torch.long)  # Ensure LongTensor
         attention_mask = inputs['attention_mask']
 
         # Forward pass
@@ -167,11 +167,14 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
 
         # 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)
             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)
             attributions, _ = ig.attribute(
                 inputs=input_ids,
                 additional_forward_args=(attention_mask,),
@@ -192,7 +195,7 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
             plt.savefig(buf, format='png', bbox_inches='tight')
             buf.seek(0)
             img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-            html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attribution Plot" style="max-width:100?%"/>')
+            html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attribution Plot" style="max-width:100%;"/>')
             plt.close()
         except Exception as e:
             logger.warning(f"Integrated Gradients failed: {e}")
@@ -257,7 +260,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.")