bertinterp.py

# -*- coding: utf-8 -*-
"""BertInterp.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1C0wG5_fRe-lH8BN8B3_4I5nbZ9cSK1HU

This notebook explores using BERT for text classification.  Before starting, change the runtime to GPU: Runtime > Change runtime type > Hardware accelerator: GPU.
"""

from google.colab import drive
drive.mount('/content/drive')

!pip install transformers

!pip install captum

from transformers import BertModel, BertTokenizer
import torch
from tqdm import tqdm
import torch.nn as nn
import numpy as np
import random
import time
import re

from captum.attr import LayerIntegratedGradients, visualization

device = torch.device("cuda")
print("Running on {}".format(device))

def read_labels(filename):
    labels={}
    with open(filename) as file:
        for line in file:
            cols = line.split("\t")
            label = cols[0]
            if label not in labels:
                labels[label]=len(labels)
    return labels

def read_data(filename, tokenizer, labels, max_data_points=None):
    data = []
    data_labels = []
    with open(filename) as file:
        for line in file:
            cols = line.split("\t")
            label = cols[0]

            sample_toks = tokenizer.tokenize(cols[1])
            text = tokenizer.convert_tokens_to_string(sample_toks)
            
            data.append(text)
            data_labels.append(labels[label])


    # shuffle the data
    tmp = list(zip(data, data_labels))
    random.shuffle(tmp)
    data, data_labels = zip(*tmp)

    if max_data_points is None:
        return data, data_labels

    return data[:max_data_points], data_labels[:max_data_points]

class BERTClassifier(nn.Module):

    
    def __init__(self, params):
        super().__init__()
    
        self.model_name=params["model_name"]
        self.base_model_name=params["base_model"]
        self.tokenizer = BertTokenizer.from_pretrained(self.base_model_name, do_lower_case=params["doLowerCase"], do_basic_tokenize=False)
        self.bert = BertModel.from_pretrained(self.base_model_name)
        # make sure it's loading on GPU
        self.bert = self.bert.to(device)

        self.bert.resize_token_embeddings(len(self.tokenizer))
        
        self.num_labels = params["label_length"]
        # make sure it's loading on GPU
        self.fc = nn.Linear(params["embedding_size"], self.num_labels).to(device)

    def get_batches(self, all_x, all_y, batch_size=2, max_toks=512):
            
        """ Get batches for input x, y data, with data tokenized according to the BERT tokenizer 
      (and limited to a maximum number of WordPiece tokens """

        batches_x=[]
        batches_y=[]
        
        for i in range(0, len(all_x), batch_size):

            current_batch=[]

            x=all_x[i:i+batch_size]

            batch_x = self.tokenizer(x, padding=True, truncation=True, return_tensors="pt", max_length=max_toks)
            batch_y=all_y[i:i+batch_size]

            batches_x.append(batch_x.to(device))
            batches_y.append(torch.LongTensor(batch_y).to(device))
            
        return batches_x, batches_y
  

    def forward(self, inputs, token_type_ids, attention_mask): 

        # double down on ensuring everything is on GPU (cuda:0)
        if inputs.device != device:
          inputs.to(device)
        if token_type_ids.device != device:
          token_type_ids.to(device)
        if attention_mask.device != device:
          attention_mask.to(device)

        bert_output=self.bert(input_ids=inputs, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True)

        bert_hidden_states = bert_output['hidden_states']

        # We're going to represent an entire document just by its [CLS] embedding (at position 0)
        out = bert_hidden_states[-1][:,0,:]

        out = self.fc(out)

        return out.squeeze()

def interpret(x, y, model, rev_labels):
    
    ''' https://captum.ai/tutorials/IMDB_TorchText_Interpret '''
    
    model.eval()
    _, maxlen=x["input_ids"].shape
    
    # baseline is uninformative sequence of padding tokens
    baseline=torch.LongTensor([[0]*maxlen]).to(device)

    y_preds=model.forward(x["input_ids"], x["token_type_ids"], x["attention_mask"])
    
    y_preds=torch.nn.functional.softmax(y_preds, dim=1)
    y_preds=y_preds.cpu().data.numpy()
    preds=[]
    for y_pred in y_preds:
        prediction=np.argmax(y_pred)
        preds.append(prediction)
    
    # we'll get our attributions with respect to target class #0 (detail!)
    target_class=0
    # be sure this is on gpu...
    bert_model.bert.embeddings.to(device)
    ig = LayerIntegratedGradients(bert_model, bert_model.bert.embeddings)
    attributions, delta = ig.attribute(x["input_ids"], baseline, target=0, return_convergence_delta=True,additional_forward_args=(x["token_type_ids"], x["attention_mask"]))
    attributions = attributions.sum(dim=2).squeeze(0)
    attributions = attributions / torch.norm(attributions)
    attributions = attributions.cpu().numpy()
    # container for tokens (for display)
    orig=[]
    # container for tokens and attr vals (for return)
    orig_and_val = []
    
    # don't think this ever gets hit but can keep for future flex
    def filter_word(word):
        word=re.sub("__NEWLINE__", "", word)
        word=re.sub("\[PAD\]", "", word)
        word=re.sub("^#+", "", word)
        return word
    
    # for each batch of input ids
    # extract the original words
    for idx, sent in enumerate(x["input_ids"]):
        toks=[]
        for word in sent:
            toks.append( filter_word(''.join(  bert_model.tokenizer.decode(word).split(" ")  )  ))
        orig.append(toks)
    

    for l in orig:
      # for each original word
      for i in range(len(l)):
        # take the sum of its attention values (the entire column of attributions at that index)

        orig_and_val.append([l[i], attributions[:, i].sum()])
    
    orig_and_val = list(zip(orig_and_val))
    
    y=y.cpu().data.numpy()
    records=[]
    # do-viz
    for idx, pred in enumerate(preds):
        records.append(visualization.VisualizationDataRecord(
                                    attributions[idx],
                                    y_preds[idx][1],
                                    rev_labels[preds[idx]],
                                    rev_labels[y[idx]],
                                    rev_labels[target_class],
                                    attributions[idx].sum(),
                                    orig[idx],
                                    delta))
    # note that legend is hard-coded in their impl
    # https://github.com/pytorch/captum/blob/master/captum/attr/_utils/visualization.py#L551-L566
    visualization.visualize_text(records, legend=True)

    return orig_and_val

labels=read_labels("/content/drive/MyDrive/ANLP21/details/train.tsv")
print(labels)
assert len(labels) == 2

rev_labels = {0: 'detail', 1: 'not_detail'}

bert_model_name="/content/drive/MyDrive/ANLP21/details/models/details-uncased_L-8_H-512_A-8.pt"
doLowerCase=False
base_model="google/bert_uncased_L-8_H-512_A-8"
bert_model = BERTClassifier(params={"base_model": base_model, "doLowerCase": doLowerCase, "model_name": bert_model_name, "embedding_size":512, "label_length": len(labels)})
bert_model.load_state_dict(torch.load(bert_model_name))

dev_x, dev_y=read_data("/content/drive/MyDrive/ANLP21/details/dev.tsv", bert_model.tokenizer, labels)

dev_batch_x, dev_batch_y = bert_model.get_batches(dev_x, dev_y)

print(len(dev_batch_x)) # 36 total, groups of 2.

orig_and_val = []

for i in range(len(dev_batch_x) - 1):
    orig_and_val.append(interpret(dev_batch_x[i], dev_batch_y[i], bert_model, rev_labels))

for i in range(len(orig_and_val)):
  filename = "dev_" + str(i) + ".txt"
  with open("/content/drive/MyDrive/ANLP21/details/" + filename, 'w') as f:
    print(orig_and_val[i], file=f)

test_x, test_y=read_data("/content/drive/MyDrive/ANLP21/details/test.tsv", bert_model.tokenizer, labels)

test_batch_x, test_batch_y = bert_model.get_batches(test_x, test_y)

orig_and_val_test = []

for i in range(len(test_batch_x) - 1):
    orig_and_val_test.append(interpret(test_batch_x[i], test_batch_y[i], bert_model, rev_labels))

for i in range(len(orig_and_val_test)):
  filename = "test_" + str(i) + ".txt"
  with open("/content/drive/MyDrive/ANLP21/details/" + filename, 'w') as f:
    print(orig_and_val_test[i], file=f)