Create generic.py

generic.py

@@ -0,0 +1,58 @@
+import os
+from transformers import BertTokenizer
+from BERT_explainability.modules.BERT.ExplanationGenerator import Generator
+from BERT_explainability.modules.BERT.BertForSequenceClassification import BertForSequenceClassification
+from transformers import BertTokenizer
+from BERT_explainability.modules.BERT.ExplanationGenerator import Generator
+from transformers import AutoTokenizer
+from captum.attr import visualization
+import spacy
+import torch
+from IPython.display import Image, HTML, display
+import sys
+
+sys.path.append('BERT')
+
+from sequenceoutput.modeling_output import SequenceClassifierOutput
+
+model = BertForSequenceClassification.from_pretrained("./BERT_weight")
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained("./BERT_weight")
+# initialize the explanations generator
+explanations = Generator(model)
+
+classifications = ["NEGATIVE", "POSITIVE"]
+true_class = 1
+
+
+def generate_visual(text_batch, target_class):
+    encoding = tokenizer(text_batch, return_tensors='pt')
+    input_ids = encoding['input_ids']
+    attention_mask = encoding['attention_mask']
+    expl = \
+        explanations.generate_LRP(input_ids=input_ids, attention_mask=attention_mask, start_layer=11,
+                                  index=target_class)[0]
+    expl = (expl - expl.min()) / (expl.max() - expl.min())
+    output = torch.nn.functional.softmax(model(input_ids=input_ids, attention_mask=attention_mask)[0], dim=-1)
+    classification = output.argmax(dim=-1).item()
+    class_name = classifications[target_class]
+    if class_name == "NEGATIVE":
+        expl *= (-1)
+    token_importance = {}
+    tokens = tokenizer.convert_ids_to_tokens(input_ids.flatten())
+    for i in range(len(tokens)):
+        token_importance[tokens[i]] = expl[i].item()
+    vis_data_records = [visualization.VisualizationDataRecord(
+        expl,
+        output[0][classification],
+        classification,
+        true_class,
+        true_class,
+        1,
+        tokens,
+        1)]
+
+    html1 = visualization.visualize_text(vis_data_records)
+    return token_importance, html1
+
+