Add code

Author: soumyasanyal

README.md

@@ -0,0 +1,32 @@
+# DIG Code
+Code for "Discretized Integrated Gradients for Explaining Language Models"
+
+<h2 align="center">
+  Overview of variants of DIG
+  <img align="center"  src="./overview.png" alt="...">
+</h2>
+**Overview of paths used in DIG and IG**. w is the word being attributed. The gray region is the neighborhood of w. Green line depicts the straight-line path from w to w' used by IG and the green squares are the corresponding interpolation points. **Left**: In DIG-Greedy, we first  monotonize each word in the neighborhood (red arrow). Then the word closest to its corresponding monotonic point is selected as the anchor (blue line to w_5 since the red arrow of w_5 has the shortest magnitude). **Right**: In DIG-MaxCount we first count the number of monotonic dimensions for each word in the neighborhood (shown in [.] above). Then, the word with the highest number of monotonic dimensions is selected as the anchor word (blue line to w_4), followed by changing the non-monotonic dimensions of w_4 (red line to c). Repeating this step gives the zigzag blue path. Finally, the red stars are the interpolated points used by our method. Please refer to the paper for more details.
+
+### Dependencies
+
+- Dependencies can be installed using `requirements.txt`.
+
+### Evaluating DIG:
+
+- Install all the requirements from `requirements.txt.`
+
+- Execute `./setup.sh` for setting up the folder hierarchy for experiments.
+
+- Commands for reproducing the reported results on DistilBERT fine-tuned on SST2:
+
+  ```shell
+  # Generate the KNN graph
+  python knn.py -dataset sst2 -nn distilbert
+  
+  # DIG (strategy: Greedy)
+  python main.py -dataset sst2 -nn distilbert -strategy greedy
+
+  # DIG (strategy: MaxCount)
+  python main.py -dataset sst2 -nn distilbert -strategy maxcount
+  ```
+  Similarly, commands can be changed for other settings.

attributions.py

@@ -0,0 +1,13 @@
+import torch
+from dig import DiscretetizedIntegratedGradients
+
+def summarize_attributions(attributions):
+	attributions = attributions.sum(dim=-1).squeeze(0)
+	attributions = attributions / torch.norm(attributions)
+	return attributions
+
+def run_dig_explanation(dig_func, all_input_embed, position_embed, type_embed, attention_mask, steps):
+	attributions 		= dig_func.attribute(scaled_features=all_input_embed, additional_forward_args=(attention_mask, position_embed, type_embed), n_steps=steps)
+	attributions_word	= summarize_attributions(attributions)
+
+	return attributions_word

bert_helper.py

@@ -0,0 +1,121 @@
+import torch, sys, pickle
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+model, tokenizer = None, None
+
+def nn_init(device, dataset, returns=False):
+	global model, tokenizer
+	if dataset == 'sst2':
+		tokenizer	= AutoTokenizer.from_pretrained('textattack/bert-base-uncased-SST-2')
+		model		= AutoModelForSequenceClassification.from_pretrained('textattack/bert-base-uncased-SST-2', return_dict=False)
+	elif dataset == 'imdb':
+		tokenizer	= AutoTokenizer.from_pretrained('textattack/bert-base-uncased-imdb')
+		model		= AutoModelForSequenceClassification.from_pretrained('textattack/bert-base-uncased-imdb', return_dict=False)
+	elif dataset == 'rotten':
+		tokenizer	= AutoTokenizer.from_pretrained('textattack/bert-base-uncased-rotten-tomatoes')
+		model		= AutoModelForSequenceClassification.from_pretrained('textattack/bert-base-uncased-rotten-tomatoes', return_dict=False)
+
+	model.to(device)
+	model.eval()
+	model.zero_grad()
+
+	if returns:
+		return model, tokenizer
+
+def move_to_device(device):
+	global model
+	model.to(device)
+
+def predict(model, inputs_embeds, attention_mask=None):
+	return model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)[0]
+
+def nn_forward_func(input_embed, attention_mask=None, position_embed=None, type_embed=None, return_all_logits=False):
+	global model
+	embeds	= input_embed + position_embed + type_embed
+	embeds	= model.bert.embeddings.dropout(model.bert.embeddings.LayerNorm(embeds))
+	pred	= predict(model, embeds, attention_mask=attention_mask)
+	if return_all_logits:
+		return pred
+	else:
+		return pred.max(1).values
+
+def load_mappings(dataset, knn_nbrs=500):
+	with open(f'processed/knns/bert_{dataset}_{knn_nbrs}.pkl', 'rb') as f:
+		[word_idx_map, word_features, adj] = pickle.load(f)
+	word_idx_map	= dict(word_idx_map)
+
+	return word_idx_map, word_features, adj
+
+def construct_input_ref_pair(tokenizer, text, ref_token_id, sep_token_id, cls_token_id, device):
+	text_ids		= tokenizer.encode(text, add_special_tokens=False, truncation=True, max_length=tokenizer.max_len_single_sentence)
+	input_ids		= [cls_token_id] + text_ids + [sep_token_id]	# construct input token ids
+	ref_input_ids	= [cls_token_id] + [ref_token_id] * len(text_ids) + [sep_token_id]	# construct reference token ids
+
+	return torch.tensor([input_ids], device=device), torch.tensor([ref_input_ids], device=device)
+
+def construct_input_ref_pos_id_pair(input_ids, device):
+	global model
+	seq_length			= input_ids.size(1)
+	position_ids 		= model.bert.embeddings.position_ids[:,0:seq_length].to(device)
+	ref_position_ids	= model.bert.embeddings.position_ids[:,0:seq_length].to(device)
+
+	return position_ids, ref_position_ids
+
+def construct_input_ref_token_type_pair(input_ids, device):
+	seq_len				= input_ids.size(1)
+	token_type_ids		= torch.tensor([[0] * seq_len], dtype=torch.long, device=device)
+	ref_token_type_ids	= torch.zeros_like(token_type_ids, dtype=torch.long, device=device)
+	return token_type_ids, ref_token_type_ids
+
+def construct_attention_mask(input_ids):
+	return torch.ones_like(input_ids)
+
+def get_word_embeddings():
+	global model
+	return model.bert.embeddings.word_embeddings.weight
+
+def construct_word_embedding(model, input_ids):
+	return model.bert.embeddings.word_embeddings(input_ids)
+
+def construct_position_embedding(model, position_ids):
+	return model.bert.embeddings.position_embeddings(position_ids)
+
+def construct_type_embedding(model, type_ids):
+	return model.bert.embeddings.token_type_embeddings(type_ids)
+
+def construct_sub_embedding(model, input_ids, ref_input_ids, position_ids, ref_position_ids, type_ids, ref_type_ids):
+	input_embeddings				= construct_word_embedding(model, input_ids)
+	ref_input_embeddings			= construct_word_embedding(model, ref_input_ids)
+	input_position_embeddings		= construct_position_embedding(model, position_ids)
+	ref_input_position_embeddings	= construct_position_embedding(model, ref_position_ids)
+	input_type_embeddings			= construct_type_embedding(model, type_ids)
+	ref_input_type_embeddings		= construct_type_embedding(model, ref_type_ids)
+
+	return 	(input_embeddings, ref_input_embeddings), \
+			(input_position_embeddings, ref_input_position_embeddings), \
+			(input_type_embeddings, ref_input_type_embeddings)
+
+def get_base_token_emb(device):
+	global model
+	return construct_word_embedding(model, torch.tensor([tokenizer.pad_token_id], device=device))
+
+def get_tokens(text_ids):
+	global tokenizer
+	return tokenizer.convert_ids_to_tokens(text_ids.squeeze())
+
+def get_inputs(text, device):
+	global model, tokenizer
+	ref_token_id = tokenizer.pad_token_id
+	sep_token_id = tokenizer.sep_token_id
+	cls_token_id = tokenizer.cls_token_id
+
+	input_ids, ref_input_ids		= construct_input_ref_pair(tokenizer, text, ref_token_id, sep_token_id, cls_token_id, device)
+	position_ids, ref_position_ids	= construct_input_ref_pos_id_pair(input_ids, device)
+	type_ids, ref_type_ids			= construct_input_ref_token_type_pair(input_ids, device)
+	attention_mask					= construct_attention_mask(input_ids)
+
+	(input_embed, ref_input_embed), (position_embed, ref_position_embed), (type_embed, ref_type_embed) = \
+				construct_sub_embedding(model, input_ids, ref_input_ids, position_ids, ref_position_ids, type_ids, ref_type_ids)
+
+	return [input_ids, ref_input_ids, input_embed, ref_input_embed, position_embed, ref_position_embed, type_embed, ref_type_embed, attention_mask]

dig.py

@@ -0,0 +1,78 @@
+import typing, sys
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from captum.log import log_usage
+from captum._utils.common import (_expand_additional_forward_args, _expand_target, _format_additional_forward_args, _format_output, _is_tuple)
+from captum._utils.typing import (BaselineType, Literal, TargetType, TensorOrTupleOfTensorsGeneric)
+from captum.attr._utils.approximation_methods import approximation_parameters
+from captum.attr._utils.attribution import GradientAttribution
+from captum.attr._utils.batching import _batch_attribution
+from captum.attr._utils.common import _format_input_baseline, _reshape_and_sum, _validate_input, _format_input
+
+
+class DiscretetizedIntegratedGradients(GradientAttribution):
+
+	def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> None:
+		GradientAttribution.__init__(self, forward_func)
+		self._multiply_by_inputs = multiply_by_inputs
+
+	@log_usage()
+	def attribute(
+		self,
+		scaled_features: Tuple[Tensor, ...],
+		target: TargetType = None,
+		additional_forward_args: Any = None,
+		n_steps: int = 50,
+		return_convergence_delta: bool = False,
+	) -> Union[
+		TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+	]:
+		is_inputs_tuple		= _is_tuple(scaled_features)
+		scaled_features_tpl	= _format_input(scaled_features)
+		attributions		= self.calculate_dig_attributions(scaled_features_tpl=scaled_features_tpl, target=target, additional_forward_args=additional_forward_args, n_steps=n_steps)
+		if return_convergence_delta:
+			assert len(scaled_features_tpl) == 1, 'More than one tuple not supported in this code!'
+			start_point, end_point = _format_input(scaled_features_tpl[0][0].unsqueeze(0)), _format_input(scaled_features_tpl[0][-1].unsqueeze(0))	# baselines, inputs (only works for one input, i.e. len(tuple) == 1)
+			# computes approximation error based on the completeness axiom
+			delta = self.compute_convergence_delta(
+				attributions,
+				start_point,
+				end_point,
+				additional_forward_args=additional_forward_args,
+				target=target,
+			)
+			return _format_output(is_inputs_tuple, attributions), delta
+
+		return _format_output(is_inputs_tuple, attributions)
+
+	def calculate_dig_attributions(
+		self,
+		scaled_features_tpl: Tuple[Tensor, ...],
+		target: TargetType = None,
+		additional_forward_args: Any = None,
+		n_steps: int = 50,
+	) -> Tuple[Tensor, ...]:
+		additional_forward_args = _format_additional_forward_args(additional_forward_args)
+		input_additional_args	= (_expand_additional_forward_args(additional_forward_args, n_steps) if additional_forward_args is not None else None)
+		expanded_target			= _expand_target(target, n_steps)
+
+		# grads: dim -> (bsz * #steps x inputs[0].shape[1:], ...)
+		grads = self.gradient_func(
+			forward_fn=self.forward_func,
+			inputs=scaled_features_tpl,
+			target_ind=expanded_target,
+			additional_forward_args=input_additional_args,
+		)
+
+		# calculate (x - x') for each interpolated point
+		shifted_inputs_tpl	= tuple(torch.cat([scaled_features[1:], scaled_features[-1].unsqueeze(0)]) for scaled_features in scaled_features_tpl)
+		steps				= tuple(shifted_inputs_tpl[i] - scaled_features_tpl[i] for i in range(len(shifted_inputs_tpl)))
+		scaled_grads		= tuple(grads[i] * steps[i] for i in range(len(grads)))
+
+		# aggregates across all steps for each tensor in the input tuple
+		attributions		= tuple(_reshape_and_sum(scaled_grad, n_steps, grad.shape[0] // n_steps, grad.shape[1:]) for (scaled_grad, grad) in zip(scaled_grads, grads))
+
+		return attributions

distilbert_helper.py

@@ -0,0 +1,107 @@
+import torch, sys, pickle
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+model, tokenizer = None, None
+
+def nn_init(device, dataset, returns=False):
+	global model, tokenizer
+	if dataset == 'sst2':
+		tokenizer	= AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+		model		= AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", return_dict=False)
+	elif dataset == 'imdb':
+		tokenizer	= AutoTokenizer.from_pretrained("textattack/distilbert-base-uncased-imdb")
+		model		= AutoModelForSequenceClassification.from_pretrained("textattack/distilbert-base-uncased-imdb", return_dict=False)
+	elif dataset == 'rotten':
+		tokenizer	= AutoTokenizer.from_pretrained("textattack/distilbert-base-uncased-rotten-tomatoes")
+		model		= AutoModelForSequenceClassification.from_pretrained("textattack/distilbert-base-uncased-rotten-tomatoes", return_dict=False)
+
+	model.to(device)
+	model.eval()
+	model.zero_grad()
+
+	if returns:
+		return model, tokenizer
+
+def move_to_device(device):
+	global model
+	model.to(device)
+
+def predict(model, inputs_embeds, attention_mask=None):
+	return model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)[0]
+
+def nn_forward_func(input_embed, attention_mask=None, position_embed=None, type_embed=None, return_all_logits=False):
+	global model
+	embeds	= input_embed + position_embed
+	embeds	= model.distilbert.embeddings.dropout(model.distilbert.embeddings.LayerNorm(embeds))
+	pred	= predict(model, embeds, attention_mask=attention_mask)
+	if return_all_logits:
+		return pred
+	else:
+		return pred.max(1).values
+
+def load_mappings(dataset, knn_nbrs=500):
+	with open(f'processed/knns/distilbert_{dataset}_{knn_nbrs}.pkl', 'rb') as f:
+		[word_idx_map, word_features, adj] = pickle.load(f)
+	word_idx_map	= dict(word_idx_map)
+
+	return word_idx_map, word_features, adj
+
+def construct_input_ref_pair(tokenizer, text, ref_token_id, sep_token_id, cls_token_id, device):
+	text_ids		= tokenizer.encode(text, add_special_tokens=False, truncation=True, max_length=tokenizer.max_len_single_sentence)
+	input_ids		= [cls_token_id] + text_ids + [sep_token_id]	# construct input token ids
+	ref_input_ids	= [cls_token_id] + [ref_token_id] * len(text_ids) + [sep_token_id]	# construct reference token ids
+
+	return torch.tensor([input_ids], device=device), torch.tensor([ref_input_ids], device=device)
+
+def construct_input_ref_pos_id_pair(input_ids, device):
+	seq_length			= input_ids.size(1)
+	position_ids		= torch.arange(seq_length, dtype=torch.long, device=device)
+	ref_position_ids	= torch.zeros(seq_length, dtype=torch.long, device=device)
+
+	position_ids		= position_ids.unsqueeze(0).expand_as(input_ids)
+	ref_position_ids	= ref_position_ids.unsqueeze(0).expand_as(input_ids)
+	return position_ids, ref_position_ids
+
+def construct_attention_mask(input_ids):
+	return torch.ones_like(input_ids)
+
+def get_word_embeddings():
+	global model
+	return model.distilbert.embeddings.word_embeddings.weight
+
+def construct_word_embedding(model, input_ids):
+	return model.distilbert.embeddings.word_embeddings(input_ids)
+
+def construct_position_embedding(model, position_ids):
+	return model.distilbert.embeddings.position_embeddings(position_ids)
+
+def construct_sub_embedding(model, input_ids, ref_input_ids, position_ids, ref_position_ids):
+	input_embeddings				= construct_word_embedding(model, input_ids)
+	ref_input_embeddings			= construct_word_embedding(model, ref_input_ids)
+	input_position_embeddings		= construct_position_embedding(model, position_ids)
+	ref_input_position_embeddings	= construct_position_embedding(model, ref_position_ids)
+
+	return (input_embeddings, ref_input_embeddings), (input_position_embeddings, ref_input_position_embeddings)
+
+def get_base_token_emb(device):
+	global model
+	return construct_word_embedding(model, torch.tensor([tokenizer.pad_token_id], device=device))
+
+def get_tokens(text_ids):
+	global tokenizer
+	return tokenizer.convert_ids_to_tokens(text_ids.squeeze())
+
+def get_inputs(text, device):
+	global model, tokenizer
+	ref_token_id = tokenizer.pad_token_id
+	sep_token_id = tokenizer.sep_token_id
+	cls_token_id = tokenizer.cls_token_id
+
+	input_ids, ref_input_ids		= construct_input_ref_pair(tokenizer, text, ref_token_id, sep_token_id, cls_token_id, device)
+	position_ids, ref_position_ids	= construct_input_ref_pos_id_pair(input_ids, device)
+	attention_mask					= construct_attention_mask(input_ids)
+
+	(input_embed, ref_input_embed), (position_embed, ref_position_embed) = construct_sub_embedding(model, input_ids, ref_input_ids, position_ids, ref_position_ids)
+
+	return [input_ids, ref_input_ids, input_embed, ref_input_embed, position_embed, ref_position_embed, None, None, attention_mask]

knn.py

@@ -0,0 +1,40 @@
+import os, sys, numpy as np, pickle, argparse
+from sklearn.neighbors import kneighbors_graph
+
+import torch
+
+
+def main(args):
+	device = torch.device("cpu")
+
+	if args.nn == 'distilbert':
+		from distilbert_helper import nn_init, get_word_embeddings, get_base_token_emb
+	elif args.nn == 'roberta':
+		from roberta_helper import nn_init, get_word_embeddings, get_base_token_emb
+	elif args.nn == 'bert':
+		from bert_helper import nn_init, get_word_embeddings, get_base_token_emb
+
+	print(f'Starting KNN computation..')
+
+	model, tokenizer	= nn_init(device, args.dataset, returns=True)
+	word_features		= get_word_embeddings().cpu().detach().numpy()
+	word_idx_map		= tokenizer.get_vocab()
+	A					= kneighbors_graph(word_features, args.nbrs, mode='distance', n_jobs=args.procs)
+
+	knn_fname = f'processed/knns/{args.nn}_{args.dataset}_{args.nbrs}.pkl'
+	with open(knn_fname, 'wb') as f:
+		pickle.dump([word_idx_map, word_features, A], f)
+
+	print(f'Written KNN data at {knn_fname}')
+
+
+if __name__ == '__main__':
+	parser = argparse.ArgumentParser(description='knn')
+	parser.add_argument('-nn',    	default='distilbert', choices=['distilbert', 'roberta', 'bert'])
+	parser.add_argument('-dataset', default='sst2', choices=['sst2', 'imdb', 'rotten'])
+	parser.add_argument('-procs',	default=40, type=int)
+	parser.add_argument('-nbrs',  	default=500, type=int)
+
+	args = parser.parse_args()
+
+	main(args)