dl-backtrace 0.0.14__py3-none-any.whl → 0.0.16.dev4__py3-none-any.whl

dl_backtrace/pytorch_backtrace/backtrace/backtrace.py

@@ -4,6 +4,9 @@ import torch.nn as nn
 from dl_backtrace.pytorch_backtrace.backtrace.utils import contrast as UC
 from dl_backtrace.pytorch_backtrace.backtrace.utils import prop as UP
 from dl_backtrace.pytorch_backtrace.backtrace.config import activation_master
+from dl_backtrace.pytorch_backtrace.backtrace.utils import helper as HP
+from dl_backtrace.pytorch_backtrace.backtrace.utils import encoder as EN
+from dl_backtrace.pytorch_backtrace.backtrace.utils import encoder_decoder as ED
 
 class Backtrace(object):
     """
@@ -11,47 +14,78 @@ class Backtrace(object):
     It takes two optional parameters: model (a neural network model) and activation_dict (a dictionary that maps layer names to activation functions).
     """
 
-    def __init__(self, model=None, activation_dict={}):
-
-        # create a tree-like structure that represents the layers of the neural network model
-        self.create_tree(model)
-
-        # create a new model (an instance of tf.keras.Model) that produces the output of each layer in the neural network.
-        self.create_model_output(model)
-
-        # create a new model (an instance of tf.keras.Model) that produces the output of each layer in the neural network.
-        self.create_every_model_output(model)
-
-        # create a layer stack that defines the order in which layers should be processed during backpropagation.
-        self.create_layer_stack()
-
-        # checks if the model is sequential or not. If it's sequential, it adds the input layer to the layer stack.
-        # identity
-
-        inp_name = 'identity'
-        self.layer_stack.append(inp_name)
-        self.model_resource[1][inp_name] = {}
-        self.model_resource[1][inp_name]["name"] = inp_name
-        self.model_resource[1][inp_name]["type"] = "input"
-        self.model_resource[1][inp_name]["parent"] = []
-        self.model_resource[1][inp_name]["child"] = None
-        self.model_resource[3].append(inp_name)
-        self.sequential = True
-        try:
+    def __init__(self, model=None, activation_dict={}, model_type=None):
+        if model_type == 'encoder':
+            self.model = model
+            self.model_type = model_type
+            
+            # create a tree-like structure for encoder model
+            self.model_resource = EN.build_encoder_tree(model)
+            
+            # create a layer stack for encoder model
+            self.create_layer_stack()
+            
+            # extract the encoder model weights
+            self.model_weights = EN.extract_encoder_weights(model)
+            
+            # # calculate the output of each submodule of the encoder model
+            # self.all_out_model = EN.create_encoder_output(model)
+            
+        elif model_type == 'encoder_decoder':
+            self.model = model
+            self.model_type = model_type
+            
+            # create a tree-like structure and layer_stack for encoder-decoder model
+            self.model_resource, self.layer_stack = ED.build_enc_dec_tree(model)
+            
+            # extract the encoder-decoder model weights
+            self.model_weights = ED.extract_encoder_decoder_weights(model)
+                       
+            # # calculate the output of each submodule of the encoder-decoder model
+            # self.all_out_model = ED.calculate_encoder_decoder_output(model)
+            
+        
+        else:
+            self.model_type = model_type
+            # create a tree-like structure that represents the layers of the neural network model
+            self.create_tree(model)
+
+            # create a new model (an instance of tf.keras.Model) that produces the output of each layer in the neural network.
+            self.create_model_output(model)
+
+            # create a new model (an instance of tf.keras.Model) that produces the output of each layer in the neural network.
+            self.create_every_model_output(model)
+
+            # create a layer stack that defines the order in which layers should be processed during backpropagation.
+            self.create_layer_stack()
+
+            # checks if the model is sequential or not. If it's sequential, it adds the input layer to the layer stack.
+            # identity
+
+            inp_name = 'identity'
+            self.layer_stack.append(inp_name)
+            self.model_resource[1][inp_name] = {}
+            self.model_resource[1][inp_name]["name"] = inp_name
+            self.model_resource[1][inp_name]["type"] = "input"
+            self.model_resource[1][inp_name]["parent"] = []
+            self.model_resource[1][inp_name]["child"] = None
+            self.model_resource[3].append(inp_name)
+            self.sequential = True
+            try:
 
-            # calls the build_activation_dict method to build a dictionary that maps layer names to activation functions.
-            # If that fails, it creates a temporary dictionary with default activation functions.
-            if len(activation_dict) == 0:
-                self.build_activation_dict(model)
-            else:
-                self.activation_dict = activation_dict
+                # calls the build_activation_dict method to build a dictionary that maps layer names to activation functions.
+                # If that fails, it creates a temporary dictionary with default activation functions.
+                if len(activation_dict) == 0:
+                    self.build_activation_dict(model)
+                else:
+                    self.activation_dict = activation_dict
 
-        except Exception as e:
-            print(e)
-            temp_dict = {}
-            for l in model.layers:
-                temp_dict[l.name] = activation_master["None"]
-            self.activation_dict = temp_dict
+            except Exception as e:
+                print(e)
+                temp_dict = {}
+                for l in model.layers:
+                    temp_dict[l.name] = activation_master["None"]
+                self.activation_dict = temp_dict 
 
     def build_activation_dict(self, model):
         model_resource = self.model_resource
@@ -321,6 +355,7 @@ class Backtrace(object):
             multiplier=100.0,
             scaler=0,
             max_unit=0,
+            predicted_token=None,
     ):
         # This method is used for evaluating layer-wise relevance based on different modes.
         if mode == "default":
@@ -330,6 +365,7 @@ class Backtrace(object):
                 multiplier=multiplier,
                 scaler=0,
                 max_unit=0,
+                predicted_token=predicted_token,
             )
             return output
         elif mode == "contrast":
@@ -342,7 +378,7 @@ class Backtrace(object):
             return output
 
     def proportional_eval(
-            self, all_out, start_wt=[], multiplier=100.0, scaler=0, max_unit=0
+            self, all_out, start_wt=[], multiplier=100.0, scaler=0, max_unit=0, predicted_token=None
     ):
         model_resource = self.model_resource
         activation_dict = self.activation_dict
@@ -350,10 +386,21 @@ class Backtrace(object):
         out_layer = model_resource[2][0]
         all_wt = {}
         if len(start_wt) == 0:
-            start_wt = UP.calculate_start_wt(all_out[out_layer])
-        all_wt[out_layer] = start_wt * multiplier
-        layer_stack = self.layer_stack
-
+            if self.model_type == 'encoder':
+                start_wt = UP.calculate_start_wt(all_out[out_layer].detach().numpy())
+                all_wt[out_layer] = start_wt * multiplier
+                layer_stack = self.layer_stack
+                all_wts = self.model_weights
+            if self.model_type == 'encoder_decoder':
+                start_wt = UP.calculate_enc_dec_start_wt(all_out[out_layer][0].detach().numpy(), predicted_token)
+                all_wt[out_layer] = start_wt * multiplier
+                layer_stack = self.layer_stack
+                all_wts = self.model_weights
+            else:
+                start_wt = UP.calculate_start_wt(all_out[out_layer])
+                all_wt[out_layer] = start_wt * multiplier
+                layer_stack = self.layer_stack
+                
         for start_layer in layer_stack:
             if model_resource[1][start_layer]["child"]:
                 child_nodes = model_resource[1][start_layer]["child"]
@@ -453,6 +500,88 @@ class Backtrace(object):
                         all_wt[start_layer], lstm_obj_f.compute_log
                     )
                     all_wt[child_nodes[0]] += temp_wt
+                    
+                elif model_resource[1][start_layer]["class"] == "Self_Attention":
+                    weights = all_wts[start_layer]
+                    self_attention_weights = HP.rename_self_attention_keys(weights)
+
+                    temp_wt = UP.calculate_wt_self_attention(
+                        all_wt[start_layer],
+                        all_out[child_nodes[0]][0].detach().numpy(),
+                        self_attention_weights,
+                    )
+                    all_wt[child_nodes[0]] += temp_wt
+                
+                elif model_resource[1][start_layer]["class"] == 'Residual':
+                    temp_wt = UP.calculate_wt_add(
+                        all_wt[start_layer],
+                        [all_out[ch].detach().numpy() for ch in child_nodes],
+                    )
+
+                    for ind, ch in enumerate(child_nodes):
+                        all_wt[ch] += temp_wt[ind]
+                
+                elif model_resource[1][start_layer]["class"] == 'Feed_Forward':
+                    weights = all_wts[start_layer]
+                    feed_forward_weights = HP.rename_feed_forward_keys(weights)
+
+                    temp_wt = UP.calculate_wt_feed_forward(
+                        all_wt[start_layer],
+                        all_out[child_nodes[0]][0].detach().numpy(),
+                        feed_forward_weights
+                    )
+                    all_wt[child_nodes[0]] += temp_wt
+                    
+                elif model_resource[1][start_layer]["class"] == "Pooler":
+                    weights = all_wts[start_layer]
+                    pooler_weights = HP.rename_pooler_keys(weights)
+
+                    temp_wt = UP.calculate_wt_pooler(
+                        all_wt[start_layer],
+                        all_out[child_nodes[0]][0].detach().numpy(),
+                        pooler_weights
+                    )
+                    all_wt[child_nodes[0]] += temp_wt
+                    
+                elif model_resource[1][start_layer]["class"] == "Classifier":
+                    weights = all_wts[start_layer]
+                    classifier_weights = HP.rename_classifier_keys(weights)
+
+                    temp_wt = UP.calculate_wt_classifier(
+                        all_wt[start_layer],
+                        all_out[child_nodes[0]][0].detach().numpy(),
+                        classifier_weights
+                    )
+                    all_wt[child_nodes[0]] += temp_wt
+                    
+                elif model_resource[1][start_layer]["class"] == "LM_Head":
+                    weights = all_wts[start_layer]
+                    lm_head_weights = HP.rename_decoder_lm_head(weights)
+
+                    temp_wt = UP.calculate_wt_lm_head(
+                        all_wt[start_layer],
+                        all_out[child_nodes[0]][0].detach().numpy(),
+                        lm_head_weights
+                    )
+                    all_wt[child_nodes[0]] += temp_wt
+                    
+                elif model_resource[1][start_layer]["class"] == 'Layer_Norm':
+                    temp_wt = all_wt[start_layer]
+                    all_wt[child_nodes[0]] += temp_wt
+                
+                elif model_resource[1][start_layer]["class"] == 'Cross_Attention':
+                    weights = all_wts[start_layer]
+                    cross_attention_weights = HP.rename_cross_attention_keys(weights)
+
+                    temp_wt = UP.calculate_wt_cross_attention(
+                        all_wt[start_layer],
+                        [all_out[ch][0].detach().numpy() for ch in child_nodes],
+                        cross_attention_weights,
+                    )
+
+                    for ind, ch in enumerate(child_nodes):
+                        all_wt[ch] += temp_wt[ind]
+
                 else:
                     temp_wt = all_wt[start_layer]
                     all_wt[child_nodes[0]] += temp_wt

dl_backtrace/pytorch_backtrace/backtrace/utils/__init__.py

@@ -1,2 +1,5 @@
 from .contrast import *
 from .prop import *
+from .encoder import *
+from .encoder_decoder import *
+from .helper import *

dl_backtrace/pytorch_backtrace/backtrace/utils/encoder.py

@@ -0,0 +1,183 @@
+import torch
+
+
+def build_encoder_tree(model, root='bert'):
+    # Initialize the tree structure
+    ltree = {}
+    layer_tree = {}
+    inputs = []
+    outputs = []
+    intermediates = []
+
+    # Base component setup
+    def add_component(tree, name, component, child=None):
+        tree[name] = {
+            'name': name,
+            'class': component if type(component).__name__ == 'str' else type(component).__name__,
+            'type': str(type(component)),
+            'parent': None,
+            'child': None
+        }
+
+        if isinstance(child, list):
+            tree[name]['child'] = child
+        elif isinstance(child, str):
+            tree[name]['child'] = [child]
+
+        if tree[name]['class'] == 'list':
+            tree[name]['class'] = [type(item).__name__ for item in component]
+            tree[name]['type'] = [str(type(item)) for item in component]
+
+        # Keep track of component type in a separate dictionary
+        layer_tree[name] = component if type(component).__name__ == 'str' else tree[name]['type']
+
+        # Link the parent to its children
+        if isinstance(child, list):
+            for ch in child:
+                if ch in tree:
+                    tree[ch]['parent'] = [name]
+
+        elif isinstance(child, str):
+            if child in tree:
+                tree[child]['parent'] = [name]
+
+        return tree[name]
+
+    # Add root and embeddings component
+    embeddings = add_component(ltree, 'embeddings', 'Embeddings', child=None)
+
+    # Add encoder layers dynamically
+    current_child = 'embeddings'
+    for i, layer in enumerate(model.bert.encoder.layer):
+        attention = add_component(ltree, f'attention_{i}', 'Self_Attention', child=current_child)
+        add_and_layer_norm_0 = add_component(ltree, f'add_and_layer_norm_{i}_0', 'Residual', child=[f'attention_{i}', current_child])
+        feed_forward = add_component(ltree, f'feed_forward_{i}', 'Feed_Forward', child=f'add_and_layer_norm_{i}_0')
+        add_and_layer_norm_1 = add_component(ltree, f'add_and_layer_norm_{i}_1', 'Residual', child=[f'feed_forward_{i}', f'add_and_layer_norm_{i}_0'])
+        current_child = f'add_and_layer_norm_{i}_1'  # Update current_child to link this layer's output to the next layer's input
+
+    # Optionally add pooler layer if present
+    if hasattr(model.bert, 'pooler'):
+        pooler = add_component(ltree, 'pooler', 'Pooler', child=current_child)
+        current_child = 'pooler'
+
+    if hasattr(model, 'classifier'):
+        classifier = add_component(ltree, 'classifier', 'Classifier', child=current_child)
+        current_child = 'classifier'
+
+    # Classify components
+    for name, component in ltree.items():
+        if component['parent'] is None:
+            outputs.append(component['name'])
+        elif component['child'] is None:
+            inputs.append(component['name'])
+        else:
+            intermediates.append(component['name'])
+
+    model_resource = (layer_tree, ltree, outputs, inputs)
+    return model_resource
+
+
+def extract_encoder_weights(model):
+    # Initialize a dictionary to hold the weights
+    weights_dict = {
+        'embeddings': {},
+        'pooler': {},
+        'dropout': {},
+        'classifier': {}
+    }
+
+    for i in range(model.config.num_hidden_layers):
+        weights_dict[f'attention_{i}'] = {}
+        weights_dict[f'add_and_layer_norm_{i}_0'] = {}
+        weights_dict[f'feed_forward_{i}'] = {}
+        weights_dict[f'add_and_layer_norm_{i}_1'] = {}
+
+    # Extract the model's parameters and organize them into the dictionary
+    for name, param in model.bert.named_parameters():
+        if 'embeddings' in name:
+            weights_dict['embeddings'][name] = param.data.numpy()
+        elif 'encoder.layer' in name:
+            layer = name.split('.')[2]
+            submodule = name.split('.')[3]
+            if 'attention' in submodule and 'LayerNorm' not in name:
+                weights_dict[f'attention_{layer}'][name] = param.data.numpy()
+            elif 'attention.output.LayerNorm' in name:
+                weights_dict[f'add_and_layer_norm_{layer}_0'][name] = param.data.numpy()
+            elif 'intermediate' in submodule:
+                weights_dict[f'feed_forward_{layer}'][name] = param.data.numpy()
+            elif 'output' in submodule and 'LayerNorm' not in name:
+                weights_dict[f'feed_forward_{layer}'][name] = param.data.numpy()
+            elif 'output.LayerNorm' in name:
+                weights_dict[f'add_and_layer_norm_{layer}_1'][name] = param.data.numpy()
+        elif 'pooler' in name:
+            weights_dict['pooler'][name] = param.data.numpy()
+
+    for name, param in model.named_parameters():
+        if 'dropout' in name:
+            weights_dict['dropout'][name] = param.data.numpy()
+        elif 'classifier' in name:
+             weights_dict['classifier'][name] = param.data.numpy()
+
+    return weights_dict
+
+
+def create_encoder_output(model, input_ids=None, attention_mask=None, token_type_ids=None):
+    all_layer_outputs = {}
+
+    # Embeddings
+    embedding_output = model.bert.embeddings(input_ids=input_ids, token_type_ids=token_type_ids)
+    all_layer_outputs['embeddings'] = embedding_output
+
+    # iterate over each layer
+    hidden_states = embedding_output
+
+    for i, layer_module in enumerate(model.bert.encoder.layer):
+        # code here
+        with torch.no_grad():
+            # Self-Attention and attention output
+            attention_output = layer_module.attention.self(
+                hidden_states,
+                attention_mask=attention_mask,
+            )[0]
+
+            # Add + Layer Norm after attention
+            attention_output = layer_module.attention.output.dense(attention_output)
+            attention_output = layer_module.attention.output.dropout(attention_output)
+            residual_attention_output = attention_output + hidden_states
+            attention_output_norm = layer_module.attention.output.LayerNorm(residual_attention_output)
+
+            # Feed Forward (Intermediate)
+            intermediate_output = layer_module.intermediate(attention_output_norm)
+
+            # Feed Forward Output
+            feed_forward_output = layer_module.output.dense(intermediate_output)
+            feed_forward_output = layer_module.output.dropout(feed_forward_output)
+            residual_feed_forward_output = feed_forward_output + attention_output_norm
+            feed_forward_output_norm = layer_module.output.LayerNorm(residual_feed_forward_output)
+
+
+        # Save outputs add_and_layer_norm_0_0
+        all_layer_outputs[f'attention_{i}'] = attention_output
+        all_layer_outputs[f'add_and_layer_norm_{i}_0'] = attention_output_norm
+        all_layer_outputs[f'feed_forward_{i}'] = feed_forward_output
+        all_layer_outputs[f'add_and_layer_norm_{i}_1'] = feed_forward_output_norm
+
+         # Update hidden states for the next layer
+        hidden_states = feed_forward_output_norm
+
+    # Pooler
+    if hasattr(model.bert, 'pooler'):
+        pooled_output = model.bert.pooler(hidden_states)
+        all_layer_outputs['pooler'] = pooled_output
+
+    if hasattr(model, 'dropout'):
+        dropout_output = model.dropout(pooled_output)
+        all_layer_outputs['dropout'] = dropout_output
+
+    # Classifier
+    if hasattr(model, 'classifier'):
+        classifier = model.classifier(dropout_output)
+        softmax_output = torch.nn.functional.softmax(classifier)
+        all_layer_outputs['classifier'] = softmax_output
+
+    return all_layer_outputs