#!/usr/bin/env python3 from collections import defaultdict import torch from pytext.models.embeddings.dict_embedding import DictEmbedding from pytext.models.embeddings.word_embedding import WordEmbedding from pytext.models.model import EmbeddingBase, EmbeddingList class PyTextInterpretableEmbedding(EmbeddingBase): r""" In PyText DocNN models we need a way to access word embedding layers, generate the embeddings and subtract the baseline. To do so, we separate embedding layers from the model, compute the embeddings separately and do all operations needed outside of the model. The original embedding layer is being replaced by `PyTextInterpretableEmbedding` layer which passes precomputed embedding vectors to lower layers. """ def __init__(self, embeddings) -> None: self.embedding_dims = [embedding.embedding_dim for embedding in embeddings] super().__init__(sum(self.embedding_dims)) self.embeddings = embeddings def forward(self, input): r""" The forward pass of embedding layer. This can be for the text or any type of embedding. Args input: Input embeddings tensor Return output: Output tensor is the same as input. It passes through the embedding tensors to lower layers without any modifications """ return input def get_attribution_map(self, attributions): r""" After attribution scores are computed for an input embedding vector we need to split it up into attribution sub tensors for each feature type: word, dict and other types TODO: we can potentally also output tuples of attributions. This might be a better option. We'll work on this in a separate diff. Args attributions: A tensor that contains attribution values for each input field. It usually has the same dimensions as the input tensor Return attribution_map: A dictionary of feature_type and attribution values """ begin = 0 attribution_map = defaultdict() for embedding, embedding_size in zip(self.embeddings, self.embedding_dims): end = begin + embedding_size if isinstance(embedding, WordEmbedding): attribution_map["word"] = attributions[:, :, begin:end] elif isinstance(embedding, DictEmbedding): attribution_map["dict"] = attributions[:, :, begin:end] else: raise NotImplementedError( "Currently only word and dict " "embeddings are supported" ) begin = end return attribution_map class BaselineGenerator: r""" This is an example input baseline generator for DocNN model which uses word and dict features. """ PAD = "" def __init__(self, model, data_handler, device) -> None: self.model = model self.data_handler = data_handler if "dict_feat" in data_handler.features: self.vocab_dict = data_handler.features["dict_feat"].vocab if "word_feat" in data_handler.features: self.vocab_word = data_handler.features["word_feat"].vocab self.baseline_single_word_feature = self._generate_baseline_single_word_feature( device ) self.baseline_single_dict_feature = self._generate_baseline_single_dict_feature( device ) def generate_baseline(self, integ_grads_embeddings, seq_length): r""" Generates baseline for input word and dict features. In the future we will extend it to support char and other features as well. This baseline is entirely based on the `` token. Args integ_grads_embeddings: A reference to integrated gradients embedding layer seq_length: The length of each sequence which depends on batch size Return baseline: A tuple of feature baselines Each feature type has a corresponding baseline tensor in the tuple. Currently only Dict and Word feature types are supported """ baseline = [] for embedding in integ_grads_embeddings.embeddings: if isinstance(embedding, WordEmbedding): baseline.append(self._generate_word_baseline(seq_length)) elif isinstance(embedding, DictEmbedding): baseline.append(self._generate_dict_baseline(seq_length)) else: raise NotImplementedError( "Currently only word and dict " "embeddings are supported" ) return tuple(baseline) def _generate_baseline_single_word_feature(self, device): return ( torch.tensor( [self.vocab_word.stoi[self.PAD] if hasattr(self, "vocab_word") else 0] ) .unsqueeze(0) .to(device) ) def _generate_baseline_single_dict_feature(self, device): r"""Generate dict features based on Assistant's case study by using sia_transformer: fbcode/assistant/sia/transformer/sia_transformer.py sia_transformer generates dict features in a special gazetter format See `fbsource/fbcode/pytext/models/embeddings/dict_embedding.py` It generates word dict feature embeddings for each word token. The output of SIATransformer after running it on `` token looks as following: OutputRecord(tokens=['<', 'pad', '>'], token_ranges=[(0, 1), (1, 4), (4, 5)], gazetteer_feats=['', '', ''], gazetteer_feat_lengths=[1, 1, 1], gazetteer_feat_weights=[0.0, 0.0, 0.0], characters=[['<', '', ''], ['p', 'a', 'd'], ['>', '', '']], pretrained_token_embedding=[ ], dense_feats=None) """ gazetteer_feats = [self.PAD, self.PAD, self.PAD] gazetteer_feat_lengths = [1, 1, 1] gazetteer_feat_weights = [0.0, 0.0, 0.0] gazetteer_feat_id = ( torch.tensor( [ self.vocab_dict.stoi[gazetteer_feat] if hasattr(self, "vocab_dict") else 0 for gazetteer_feat in gazetteer_feats ] ) .unsqueeze(0) .to(device) ) gazetteer_feat_weights = ( torch.tensor(gazetteer_feat_weights).unsqueeze(0).to(device) ) gazetteer_feat_lengths = ( torch.tensor(gazetteer_feat_lengths).to(device).view(1, -1)[:, 1] ) return (gazetteer_feat_id, gazetteer_feat_weights, gazetteer_feat_lengths) def _generate_word_baseline(self, seq_length): return self.baseline_single_word_feature.repeat(1, seq_length) def _generate_dict_baseline(self, seq_length): return ( self.baseline_single_dict_feature[0].repeat(1, seq_length), self.baseline_single_dict_feature[1].repeat(1, seq_length), self.baseline_single_dict_feature[2].repeat(1, seq_length), ) def configure_task_integ_grads_embeddings(task): r""" Wraps Pytext's DocNN model embedding with `IntegratedGradientsEmbedding` for a given input task. IntegratedGradientsEmbedding allows to perform baseline related operations Args task: DocNN task reference Returns integrated_gradients_embedding_lst: The embedding layer which contains IntegratedGradientsEmbedding as a wrapper over the original embeddings of the model """ integrated_gradients_embedding_lst = configure_model_integ_grads_embeddings( task.model ) task.model.embedding = integrated_gradients_embedding_lst return integrated_gradients_embedding_lst[0] def configure_model_integ_grads_embeddings(model): r""" Wraps Pytext's DocNN model embedding with `IntegratedGradientsEmbedding` IntegratedGradientsEmbedding allows to perform baseline related operations Args model: a reference to DocModel Returns integrated_gradients_embedding_lst: The embedding layer which contains IntegratedGradientsEmbedding as a wrapper over the original embeddings of the model """ embeddings = model.embedding integrated_gradients_embedding = PyTextInterpretableEmbedding(embeddings) return EmbeddingList([integrated_gradients_embedding], False) def reshape_word_features(word_features): r""" Creates one-sample batch for word features for sanity check purposes Args word_features: A tensor of diemnsions #words x #embeddings Return word_features: A tensor of dimensions 1 x #words x #embeddings """ return word_features.unsqueeze(0) def reshape_dict_features( dict_feature_id_batch, dict_weight_batch, dict_seq_len_batch, seq_length, idx ): r""" Creates one-sample batch for dict features for sanity check purposes It reads and reshapes id, weight and seq_length feature arrays for given input index `idx` from the input batch Args dict_feature_id_batch: The batch tensor for ids dict_weight_matrix: The batch tensor for weights dict_seq_len_matrix: The batch tensor for sequence length seq_length: The number of tokens per sequence idx: The index of sample in the batch Return dict_feature_ids: A tensor of dimensions [ bsz x # dict feature embeddings] dict_feature_weights: [ bsz x # dict feature embeddings] dict_feature_lens: [ bsz * seq_length ] """ dict_feature_ids = dict_feature_id_batch[idx].unsqueeze(0) dict_feature_weights = dict_weight_batch[idx].unsqueeze(0) dict_feature_lens = dict_seq_len_batch[idx].unsqueeze(0) return (dict_feature_ids, dict_feature_weights, dict_feature_lens)