File size: 5,275 Bytes
d1dc1ec
 
 
 
 
9ba191d
 
 
d1dc1ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7e95840
23e5927
d1dc1ec
 
 
 
 
a471e63
d1dc1ec
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import gradio as gr

import torch
import numpy as np
from transformers import AutoTokenizer
from src.modeling.modeling_bert import BertForSequenceClassification
from src.modeling.modeling_electra import ElectraForSequenceClassification
from src.attention_rollout import AttentionRollout

import seaborn as sns
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.backends.backend_pdf

def inference(text, model):
  if model == "bert-base-uncased-cls-sst2":
    config = {
        # As of now, BERT and ELECTRA are supported. You can choose any checkpoing of these models.
        ### BERT-base
        "MODEL": "TehranNLP-org/bert-base-uncased-cls-sst2"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-mnli"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-hatexplain"
        ### BERT-large
        # "MODEL": "TehranNLP-org/bert-large-sst2"
        # "MODEL": "TehranNLP-org/bert-large-mnli"
        # "MODEL": "TehranNLP-org/bert-large-hateXplain"
        ### ELECTRA
        # "MODEL": "TehranNLP-org/electra-base-sst2"
        # "MODEL": "TehranNLP-org/electra-base-mnli"
        # "MODEL": "TehranNLP-org/electra-base-hateXplain"
    }
  elif model == "bert-large-sst2":
    config = {
        # As of now, BERT and ELECTRA are supported. You can choose any checkpoing of these models.
        ### BERT-base
        #"MODEL": "TehranNLP-org/bert-base-uncased-cls-sst2"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-mnli"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-hatexplain"
        ### BERT-large
        "MODEL": "TehranNLP-org/bert-large-sst2"
        # "MODEL": "TehranNLP-org/bert-large-mnli"
        # "MODEL": "TehranNLP-org/bert-large-hateXplain"
        ### ELECTRA
        # "MODEL": "TehranNLP-org/electra-base-sst2"
        # "MODEL": "TehranNLP-org/electra-base-mnli"
        # "MODEL": "TehranNLP-org/electra-base-hateXplain"
    }
  else:
   config = {
        # As of now, BERT and ELECTRA are supported. You can choose any checkpoing of these models.
        ### BERT-base
        #"MODEL": "TehranNLP-org/bert-base-uncased-cls-sst2"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-mnli"
        # "MODEL": "TehranNLP-org/bert-base-uncased-cls-hatexplain"
        ### BERT-large
        #"MODEL": "TehranNLP-org/bert-large-sst2"
        # "MODEL": "TehranNLP-org/bert-large-mnli"
        # "MODEL": "TehranNLP-org/bert-large-hateXplain"
        ### ELECTRA
        "MODEL": "TehranNLP-org/electra-base-sst2"
        # "MODEL": "TehranNLP-org/electra-base-mnli"
        # "MODEL": "TehranNLP-org/electra-base-hateXplain"
    }
  SENTENCE = text
  
  tokenizer = AutoTokenizer.from_pretrained(config["MODEL"])
  tokenized_sentence = tokenizer.encode_plus(SENTENCE, return_tensors="pt")
  if "bert" in config["MODEL"]:
      model = BertForSequenceClassification.from_pretrained(config["MODEL"])
  elif "electra" in config["MODEL"]:
      model = ElectraForSequenceClassification.from_pretrained(config["MODEL"])
  else:
      raise Exception(f"Not implented model: {config['MODEL']}")
      
  # Extract single layer attentions
  with torch.no_grad():
      logits, attentions, norms = model(**tokenized_sentence, output_attentions=True, output_norms=True, return_dict=False)
      num_layers = len(attentions)
      norm_nenc = torch.stack([norms[i][4] for i in range(num_layers)]).squeeze().cpu().numpy()
      print("Single layer N-Enc token attribution:", norm_nenc.shape)
  
      # Aggregate and compute GlobEnc
      globenc = AttentionRollout().compute_flows([norm_nenc], output_hidden_states=True)[0]
      globenc = np.array(globenc)
      print("Aggregated N-Enc token attribution (GlobEnc):", globenc.shape)
      


  tokenized_text = tokenizer.convert_ids_to_tokens(tokenized_sentence["input_ids"][0])
  plt.figure(figsize=(14, 8))
  norm_cls = globenc[:, 0, :]
  norm_cls = np.flip(norm_cls, axis=0)
  row_sums = norm_cls.max(axis=1)
  norm_cls = norm_cls / row_sums[:, np.newaxis]
  df = pd.DataFrame(norm_cls, columns=tokenized_text, index=range(len(norm_cls), 0, -1))
  ax = sns.heatmap(df, cmap="Reds", square=True)
  bottom, top = ax.get_ylim()
  ax.set_ylim(bottom + 0.5, top - 0.5)
  plt.title("GlobEnc", fontsize=16)
  plt.ylabel("Layer", fontsize=16)
  plt.xticks(rotation = 90, fontsize=16)
  plt.yticks(fontsize=13)
  plt.gcf().subplots_adjust(bottom=0.2)
  print("logits:", logits)
  
  return plt
  
demo = gr.Blocks()

with demo:
    gr.Markdown(
    """
    ## Please check out [DecompX](https://huggingface.co/spaces/mohsenfayyaz/DecompX) for a more faithful approach.
    ## Gradio Demo for [mohsenfayyaz/GlobEnc](https://github.com/mohsenfayyaz/GlobEnc), GlobEnc: Quantifying Global Token Attribution by Incorporating the Whole Encoder Layer in Transformers
    """)
    inp = [gr.Textbox(),gr.Dropdown(choices=['bert-base-uncased-cls-sst2','bert-large-sst2','electra-base-sst2'])]
    out = gr.Plot()
    
    button = gr.Button(value="Run")
    gr.Examples([["A deep and meaningful film.", "bert-base-uncased-cls-sst2"], ["A deep and meaningful film.", "bert-large-sst2"]], inp, out, inference)
    button.click(fn=inference, 
               inputs=inp, 
               outputs=out)

demo.launch()