initial

Author: danielzuegner

.gitignore

@@ -0,0 +1,7 @@
+# Python cache
+_pycache__
+*.py[cod]
+
+# IntelliJ/Jupyter
+.idea
+.ipynb_checkpoints

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Daniel Zügner and Tobias Kirschstein, Technical University of Munich
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class AttentionType(Enum):
+    SCALED_DOT_PRODUCT = "scaled_dot_product"
+    ADDITIVE = "additive"
+    MULTIHEAD = "multihead"

code_transformer/configuration/attention.py

@@ -0,0 +1,48 @@
+from code_transformer.configuration.configuration_utils import ModelConfiguration
+
+
+class CodeTransformerLayerConfig(ModelConfiguration):
+    def __init__(self,
+                 d_model=256,
+                 nhead=8,
+                 dim_feedforward=1024,
+                 activation="gelu",
+                 dropout=0.1,
+
+                 num_relative_distances=1,
+                 use_token_distances=False,
+                 use_edge_embeddings=False,
+                 use_content_content=True,
+                 use_content_pos=True,
+                 use_pos_content=True,
+                 use_pos_pos=True, ):
+        super(CodeTransformerLayerConfig, self).__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+        self.dim_feedforward = dim_feedforward
+        self.activation = activation
+        self.dropout = dropout
+        self.num_relative_distances = num_relative_distances
+        self.use_token_distances = use_token_distances
+        self.use_edge_embeddings = use_edge_embeddings
+        self.use_content_content = use_content_content
+        self.use_content_pos = use_content_pos
+        self.use_pos_content = use_pos_content
+        self.use_pos_pos = use_pos_pos
+
+
+class CodeTransformerCoreConfig(ModelConfiguration):
+    def __init__(self,
+                 encoder_layer: CodeTransformerLayerConfig,
+                 num_layers: int,
+                 positional_encoding=None,
+                 norm=None
+                 ):
+        super(CodeTransformerCoreConfig, self).__init__()
+        if isinstance(encoder_layer, CodeTransformerLayerConfig):
+            self.encoder_layer = CodeTransformerLayerConfig(**encoder_layer)
+        else:
+            self.encoder_layer = encoder_layer
+        self.num_layers = num_layers
+        self.norm = norm
+        self.positional_encoding = positional_encoding

code_transformer/configuration/code_transformer.py

@@ -0,0 +1,39 @@
+class DotDict(dict):
+    """
+    Simple extension of Python's dict to support dot access.
+    """
+    def __init__(self, *args, **kwargs):
+        super(DotDict, self).__init__(*args, **kwargs)
+        for arg in args:
+            if isinstance(arg, dict):
+                for k, v in arg.items():
+                    self[k] = v
+
+        if kwargs:
+            for k, v in kwargs.items():
+                if isinstance(v, dict):
+                    self[k] = DotDict(**v)
+                else:
+                    self[k] = v
+
+    def __getattr__(self, attr):
+        return self[attr]
+
+    def __setattr__(self, key, value):
+        self.__setitem__(key, value)
+
+    def __setitem__(self, key, value):
+        super(DotDict, self).__setitem__(key, value)
+        self.__dict__.update({key: value})
+
+    def __delattr__(self, item):
+        self.__delitem__(item)
+
+    def __delitem__(self, key):
+        super(DotDict, self).__delitem__(key)
+        del self.__dict__[key]
+
+
+class ModelConfiguration(DotDict):
+    pass
+

code_transformer/configuration/configuration_utils.py

@@ -0,0 +1,44 @@
+from code_transformer.configuration.configuration_utils import ModelConfiguration
+
+
+class GreatTransformerConfig(ModelConfiguration):
+    def __init__(self,
+                 num_layers: int,
+                 positional_encoding=None,
+                 embed_dim=256,
+                 num_heads=8,
+                 ff_dim=1024,
+                 dropout_rate=0.1,
+                 is_encoder_decoder=False
+                 ):
+        super(GreatTransformerConfig, self).__init__()
+
+        self.num_layers = num_layers
+        self.positional_encoding = positional_encoding
+
+        self.embed_dim = embed_dim
+        self.hidden_dim = embed_dim
+        self.attention_dim = embed_dim
+        self.bias_dim = embed_dim
+        self.num_heads = num_heads
+        self.ff_dim = ff_dim
+        self.dropout_rate = dropout_rate
+        self.is_encoder_decoder = is_encoder_decoder
+
+
+class GreatEncoderConfig(ModelConfiguration):
+
+    def __init__(self,
+                 transformer_config: GreatTransformerConfig,
+                 vocab_size=32000,
+                 num_node_types=None,
+                 subtokens_per_token=5,
+                 num_languages=None
+                 ):
+        super(GreatEncoderConfig, self).__init__()
+
+        self.transformer_config = transformer_config
+        self.vocab_size = vocab_size
+        self.num_node_types = num_node_types
+        self.subtokens_per_token = subtokens_per_token
+        self.num_languages = num_languages

code_transformer/configuration/great_transformer.py

@@ -0,0 +1,64 @@
+from torch import nn
+
+from code_transformer.configuration.configuration_utils import ModelConfiguration
+from code_transformer.configuration.attention import AttentionType
+
+
+class TransformerLMDecoderConfig(ModelConfiguration):
+
+    def __init__(self,
+                 lm_encoder,  #: Union[TransformerLMEncoder, TransformerLMEncoderConfig],
+                 sos_id: int,
+                 unk_id=0,
+                 n_layers=1,
+                 decoder_dropout=0,
+                 decoder_nhead=8,
+                 decoder_dim_feedforward=2048,
+                 decoder_activation="gelu",
+                 use_teacher_forcing=False,
+                 output_subtokens_per_token=5,
+                 output_nonlinearity=None,
+                 loss_fct=nn.CrossEntropyLoss(ignore_index=-1),
+                 use_pointer_network=False,
+                 use_pointer_query_linear=False,
+                 use_pointer_query_self_attention=False,
+                 concat_query_and_pointer=True,
+                 attend_cls_token=False,
+                 pointer_attention_type=AttentionType.MULTIHEAD,
+                 target_vocab_size: int = None):
+        r"""
+        :param lm_encoder: The encoder on which the decoder should be built
+        :param sos_id: The ID of the SOS token in the underlying vocabulary. Initially the decoder sequence will be
+            populated with a single SOS token to have an input for the LSTM decoder
+        :param n_layers: How many layers the decoder should have.
+        :param decoder_dropout: Whether dropout should be applied in the decoder.
+        :param use_teacher_forcing: If set, the previous label will be fed into the decoder instead of the
+            previous prediction during training. Usually speeds up training but also introduces a gap between
+            training and evaluation.
+        :param target_vocab_size: If given, the model assumes separate vocabularies for input tokens and label tokens.
+            The value specified here indicates the output distribution domain that is to be predicted
+        """
+        super(TransformerLMDecoderConfig, self).__init__()
+        self.lm_encoder = lm_encoder
+        self.sos_id = sos_id
+        self.unk_id = unk_id
+        self.n_layers = n_layers
+        self.decoder_dropout = decoder_dropout
+        self.decoder_nhead = decoder_nhead
+        self.decoder_dim_feedforward = decoder_dim_feedforward
+        self.decoder_activation = decoder_activation
+        self.use_teacher_forcing = use_teacher_forcing
+        self.output_subtokens_per_token = output_subtokens_per_token
+        self.output_nonlinearity = output_nonlinearity
+        self.loss_fct = loss_fct
+        self.use_pointer_network = use_pointer_network
+        self.pointer_attention_type = pointer_attention_type if isinstance(pointer_attention_type,
+                                                                           AttentionType) else AttentionType(
+            pointer_attention_type)
+        self.use_pointer_query_linear = use_pointer_query_linear
+        self.use_pointer_query_self_attention = use_pointer_query_self_attention
+        self.concat_query_and_pointer = concat_query_and_pointer
+        self.attend_cls_token = attend_cls_token
+        assert not (
+                    use_pointer_query_self_attention and use_pointer_query_linear), "Cannot set both query linear and query self attention"
+        self.target_vocab_size = target_vocab_size

code_transformer/configuration/transformer_lm_decoder.py

@@ -0,0 +1,22 @@
+from code_transformer.configuration.configuration_utils import ModelConfiguration
+
+
+class TransformerLMEncoderConfig(ModelConfiguration):
+
+    def __init__(self,
+                 transformer,  #: Union[CodeTransformer, CodeTransformerCoreConfig],
+                 vocab_size=32000,
+                 num_node_types=None,
+                 num_token_types=None,
+                 subtokens_per_token=5,
+                 input_nonlinearity=None,
+                 num_languages=None):
+        super(TransformerLMEncoderConfig, self).__init__()
+
+        self.transformer = transformer
+        self.vocab_size = vocab_size
+        self.num_token_types = num_token_types
+        self.num_node_types = num_node_types
+        self.subtokens_per_token = subtokens_per_token
+        self.input_nonlinearity = input_nonlinearity
+        self.num_languages = num_languages

code_transformer/configuration/transformer_lm_encoder.py

@@ -0,0 +1,20 @@
+from code_transformer.experiments.experiment import ExperimentSetup, ex
+from code_transformer.experiments.mixins.code_summarization import CTCodeSummarizationMixin
+from code_transformer.experiments.mixins.code_trans_transformer import CodeTransformerDecoderMixin
+
+
+class CodeTransDecoderExperimentSetup(CodeTransformerDecoderMixin,
+                                      CTCodeSummarizationMixin,
+                                      ExperimentSetup):
+    pass
+
+
+@ex.automain
+def main():
+    experiment = CodeTransDecoderExperimentSetup()
+    experiment.train()
+
+
+@ex.command(unobserved=True)
+def recreate_experiment():
+    return CodeTransDecoderExperimentSetup()

code_transformer/experiments/code_transformer/code_summarization.py

@@ -0,0 +1,113 @@
+experiment_setup:
+  executable: 'code_transformer/experiments/code_transformer/code_summarization.py'
+
+data_setup:
+  language: 'python,javascript,ruby,go' # The dataset (language) to use
+  filter_language: python               # only for multi-language datasets. Let only snippets of the specified language through (fine-tuning)
+  use_validation: True
+  num_sub_tokens: 5                     # Number of sub-tokens that input tokens should be split into
+  num_subtokens_output: 6               # Number of sub-tokens of the method name to predict
+  use_only_ast: False                   # Whether to use the only-ast ablation
+  mask_all_tokens: False                # Only relevant if use_only_ast=True. Replaces all input tokens with a dummy token
+  use_no_punctuation: True              # Whether to drop punctuation tokens before feeding the snippets into the model
+  use_pointer_network: True             # Whether to use a pointer network in the decoder
+  sort_by_length: False                 # Whether to sort loaded slices by number of tokens. Useful to minimize amount of zero-padding needed
+  shuffle: False                        # Whether load order of snippets should be randomized
+  chunk_size: 32                        # Only relevant if shuffle=True and sort_by_lenght=True. Snippets will be chunked into chunks of `chunk_size`, which will then be randomly shuffled.
+
+data_transforms:
+  max_distance_mask: None               # Mask nodes if their relative distances exceed a certain treshold
+#    max_distance_mask:
+#      shortest_paths: 5
+#      ppr: -1
+#      sibling_sp: -1
+#      ancestor_sp: -1
+
+  relative_distances:                   # Which relative distances to use (have to be pre-computed in stage 2 preprocessing)
+    - ppr
+    - ancestor_sp
+    - sibling_sp
+    - shortest_paths
+
+  distance_binning:                     # Distance binning for dealing with real-valued distances
+    type: 'exponential'                 # "exponential" or "equal". Exponential binning has more diversified (smaller) bins for smaller distances
+    growth_factor: 1.3
+    n_fixed_bins: 9
+
+transfer_learning:                      # Load a pretrained model for fine-tuninig
+  use_pretrained_model: False
+  model_type: 'ct_code_summarization'
+  run_id: CT-23
+  snapshot_iteration: 10
+  cpu: False
+  freeze_encoder_layers: None           # None, "all" or {number}. If and how many encoder layers to keep constant during fine-tuning
+
+model:
+  with_cuda: True                       # Run model on GPU
+  label_smoothing: 0.1                  # Apply label smoothing to ground truth
+  lm_encoder:                           # Hyperparameters of the encoder
+    input_nonlinearity: 'tanh'
+    num_languages: None                   # only relevant for multi-language datasets. How many different languages have been fused together
+    transformer:                          # CodeTransformer hyperparameters
+      num_layers: 1
+      encoder_layer:
+        d_model: 1024                       # Internal embedding dimension
+        nhead: 8                            # Number of attention heads
+        dim_feedforward: 2048               # Dimension of feed-forward layer
+        dropout: 0.2
+        activation: 'gelu'
+        use_content_content: True           # Whether to use the content-content term in attention computation
+        use_content_pos: True               # Whether to use the content-content term in attention computation
+        use_pos_content: True               # Whether to use the content-content term in attention computation
+        use_pos_pos: True                   # Whether to use the content-content term in attention computation
+        use_token_distances: True           # Whether to also compute the simple hop-distance between the input tokens
+  lm_decoder:                           # Hyperparameters of the decoder
+    output_nonlinearity: None
+    n_layers: 1
+    decoder_dropout: 0
+    decoder_nhead: 8
+    decoder_dim_feedforward: 2048
+    decoder_activation: 'gelu'
+    use_teacher_forcing: True             # Whether to use teacher forcing during training (Label is fed into decoder instead of prediction for previous position)
+    pointer_attention_type: 'additive'    # Attention type in Pointer Network. "scaled_dot_product", "multihead" or "additive"
+    use_pointer_query_self_attention: False # Whether to use self-attention between pointer query and decoder input
+    concat_query_and_pointer: True        # Whether to also use the query-stream of the encoder output to guide the pointer query
+    attend_cls_token: False               # Whether to mask the CLS token for attention
+
+optimizer:
+  optimizer: 'Adam'
+  learning_rate: 8e-5
+  reg_scale: 3e-5
+
+  #scheduler: 'OneCycleLR'
+  #scheduler_params:
+  #  max_lr: 1e-4
+  #  steps_per_epoch: 4000 # 500000 / 128
+  #  epochs: 30
+  #  pct_start: 0.3
+
+  #scheduler: 'MultiStepLR'
+  #scheduler_params:
+  #  milestones: [1500, 5000]
+  #  gamma: 0.1
+
+training:
+  random_seed: 456
+  batch_size: 8
+  simulated_batch_size: 128         # Gradient Accumulation. After how many samples the model parameters should be updated
+  simulated_batch_size_valid: 1280  # Over how many samples validation metrics should be calculated
+  accumulate_tokens_batch: False
+  validate_every: 100               # Counted in number of parameter updates (simulated_batch_size). How often approximate evaluation should be done
+  persistent_snapshot_every: 10000  # Counted in loop iterations (batch_size). Also starts a full evaluation on the validation set to guide early stopping
+  early_stopping_patience: 20       # How often the model evaluation can be worse than the current best before training is stopped
+  max_validation_samples: 50000     # Limit number of samples for full evaluation on larger datasets to speed up training
+  metrics:                          # Which metrics to log
+    - top1_accuracy
+    - top5_accuracy
+    - non_trivial_accuracy
+    - precision
+    - recall
+    - f1_score
+    - micro_f1_score
+    - rouge_2
+    - rouge_l