From 0d0c73a5a0d3996ec40090cf040b9a6fb37bb744 Mon Sep 17 00:00:00 2001 From: Philip Robinson Date: Thu, 25 Mar 2021 18:26:53 -0700 Subject: [PATCH] standardize code --- .../attentional_recurrent_subtoken_decoder.py | 218 +++--- neural-model/model/embedding.py | 41 +- neural-model/model/encoder.py | 1 - neural-model/model/ensemble_model.py | 177 +++-- neural-model/model/gnn.py | 171 +++-- neural-model/model/graph_encoder.py | 610 ++++++++++------ neural-model/model/hybrid_encoder.py | 94 ++- neural-model/model/model.py | 128 ++-- neural-model/model/recurrent_decoder.py | 210 ++++-- .../model/recurrent_subtoken_decoder.py | 212 ++++-- neural-model/model/sequential_encoder.py | 168 +++-- neural-model/model/simple_decoder.py | 88 ++- neural-model/utils/ast.py | 181 +++-- neural-model/utils/code_processing.py | 90 +-- neural-model/utils/dataset.py | 514 +++++++++----- neural-model/utils/eval_debin_prediction.py | 43 +- neural-model/utils/evaluation.py | 205 ++++-- neural-model/utils/get_stat.py | 30 +- neural-model/utils/grammar.py | 10 +- neural-model/utils/graph.py | 15 +- neural-model/utils/gz_to_jsonl.py | 11 +- neural-model/utils/lexer.py | 80 ++- neural-model/utils/nn_util.py | 58 +- neural-model/utils/preprocess.py | 242 ++++--- neural-model/utils/sequential_preprocess.py | 36 +- neural-model/utils/subsample.py | 39 +- neural-model/utils/util.py | 21 +- neural-model/utils/vocab.py | 232 ++++--- .../attentional_recurrent_subtoken_decoder.py | 232 ++++--- prediction-plugin/model/embedding.py | 69 +- prediction-plugin/model/ensemble_model.py | 319 +++++++++ prediction-plugin/model/gnn.py | 201 +++--- prediction-plugin/model/graph_encoder.py | 653 +++++++++--------- prediction-plugin/model/hybrid_encoder.py | 90 ++- prediction-plugin/model/model.py | 142 ++-- prediction-plugin/model/recurrent_decoder.py | 351 ++++++++++ .../model/recurrent_subtoken_decoder.py | 212 ++++-- prediction-plugin/model/sequential_encoder.py | 168 +++-- prediction-plugin/model/simple_decoder.py | 135 ++++ prediction-plugin/utils/ast.py | 173 ++--- prediction-plugin/utils/code_processing.py | 93 ++- prediction-plugin/utils/dataset.py | 550 ++++++++------- .../utils/eval_debin_prediction.py | 95 +++ prediction-plugin/utils/evaluation.py | 208 +++--- prediction-plugin/utils/get_stat.py | 59 ++ prediction-plugin/utils/grammar.py | 14 +- prediction-plugin/utils/graph.py | 14 +- prediction-plugin/utils/gz_to_jsonl.py | 17 + prediction-plugin/utils/lexer.py | 74 +- prediction-plugin/utils/nn_util.py | 64 +- prediction-plugin/utils/preprocess.py | 181 +++-- .../utils/sequential_preprocess.py | 78 +++ prediction-plugin/utils/subsample.py | 51 ++ prediction-plugin/utils/util.py | 19 +- prediction-plugin/utils/vocab.py | 247 +++---- 55 files changed, 5386 insertions(+), 3048 deletions(-) create mode 100644 prediction-plugin/model/ensemble_model.py create mode 100644 prediction-plugin/model/recurrent_decoder.py create mode 100644 prediction-plugin/model/simple_decoder.py create mode 100644 prediction-plugin/utils/eval_debin_prediction.py create mode 100644 prediction-plugin/utils/get_stat.py create mode 100644 prediction-plugin/utils/gz_to_jsonl.py create mode 100644 prediction-plugin/utils/sequential_preprocess.py create mode 100644 prediction-plugin/utils/subsample.py diff --git a/neural-model/model/attentional_recurrent_subtoken_decoder.py b/neural-model/model/attentional_recurrent_subtoken_decoder.py index 7243a8d..8a4304c 100644 --- a/neural-model/model/attentional_recurrent_subtoken_decoder.py +++ b/neural-model/model/attentional_recurrent_subtoken_decoder.py @@ -1,37 +1,48 @@ import sys -from collections import namedtuple, OrderedDict -from typing import Dict, List, Any +from collections import OrderedDict, namedtuple +from typing import Any, Dict, List import torch -from torch import nn as nn - from model.decoder import Decoder from model.encoder import Encoder -from utils import util, nn_util +from model.recurrent_subtoken_decoder import RecurrentSubtokenDecoder +from torch import nn as nn +from utils import nn_util, util from utils.ast import AbstractSyntaxTree from utils.dataset import Example -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN, END_OF_VARIABLE_TOKEN - -from model.recurrent_subtoken_decoder import RecurrentSubtokenDecoder -from utils.vocab import Vocab +from utils.vocab import END_OF_VARIABLE_TOKEN, SAME_VARIABLE_TOKEN, Vocab class AttentionalRecurrentSubtokenDecoder(RecurrentSubtokenDecoder): - def __init__(self, variable_encoding_size: int, context_encoding_size: int, hidden_size: int, dropout: float, - tie_embed: bool, input_feed: bool, vocab: Vocab): - super(AttentionalRecurrentSubtokenDecoder, self).__init__(variable_encoding_size, hidden_size, dropout, tie_embed, input_feed, vocab) + def __init__( + self, + variable_encoding_size: int, + context_encoding_size: int, + hidden_size: int, + dropout: float, + tie_embed: bool, + input_feed: bool, + vocab: Vocab, + ): + super(AttentionalRecurrentSubtokenDecoder, self).__init__( + variable_encoding_size, hidden_size, dropout, tie_embed, input_feed, vocab + ) self.att_src_linear = nn.Linear(context_encoding_size, hidden_size, bias=False) - self.att_vec_linear = nn.Linear(context_encoding_size + hidden_size, hidden_size, bias=False) + self.att_vec_linear = nn.Linear( + context_encoding_size + hidden_size, hidden_size, bias=False + ) @classmethod def default_params(cls): params = RecurrentSubtokenDecoder.default_params() - params.update({ - 'remove_duplicates_in_prediction': True, - 'context_encoding_size': 128, - 'attention_target': 'ast_nodes' # terminal_nodes - }) + params.update( + { + "remove_duplicates_in_prediction": True, + "context_encoding_size": 128, + "attention_target": "ast_nodes", # terminal_nodes + } + ) return params @@ -39,9 +50,16 @@ def default_params(cls): def build(cls, config): params = util.update(cls.default_params(), config) - vocab = Vocab.load(params['vocab_file']) - model = cls(params['variable_encoding_size'], params['context_encoding_size'], - params['hidden_size'], params['dropout'], params['tie_embedding'], params['input_feed'], vocab) + vocab = Vocab.load(params["vocab_file"]) + model = cls( + params["variable_encoding_size"], + params["context_encoding_size"], + params["hidden_size"], + params["dropout"], + params["tie_embedding"], + params["input_feed"], + vocab, + ) model.config = params return model @@ -49,10 +67,12 @@ def build(cls, config): def rnn_step(self, x, h_tm1, context_encoding): h_t, cell_t = self.lstm_cell(x, h_tm1) - ctx_t, alpha_t = nn_util.dot_prod_attention(h_t, - context_encoding['attention_value'], - context_encoding['attention_key'], - context_encoding['attention_value_mask']) + ctx_t, alpha_t = nn_util.dot_prod_attention( + h_t, + context_encoding["attention_value"], + context_encoding["attention_key"], + context_encoding["attention_value_mask"], + ) att_t = torch.tanh(self.att_vec_linear(torch.cat([h_t, ctx_t], 1))) att_t = self.dropout(att_t) @@ -60,33 +80,39 @@ def rnn_step(self, x, h_tm1, context_encoding): return (h_t, cell_t), att_t, ctx_t def get_attention_memory(self, context_encoding): - att_tgt = self.config['attention_target'] + att_tgt = self.config["attention_target"] value, mask = self.encoder.get_attention_memory(context_encoding, att_tgt) key = self.att_src_linear(value) - return {'attention_key': key, - 'attention_value': value, - 'attention_value_mask': mask} + return { + "attention_key": key, + "attention_value": value, + "attention_value_mask": mask, + } def forward(self, src_ast_encoding, prediction_target): # prepare tensors for attention attention_memory = self.get_attention_memory(src_ast_encoding) src_ast_encoding.update(attention_memory) - return RecurrentSubtokenDecoder.forward(self, src_ast_encoding, prediction_target) + return RecurrentSubtokenDecoder.forward( + self, src_ast_encoding, prediction_target + ) def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: batch_size = len(examples) - beam_size = self.config['beam_size'] + beam_size = self.config["beam_size"] same_variable_id = self.vocab.target[SAME_VARIABLE_TOKEN] end_of_variable_id = self.vocab.target[END_OF_VARIABLE_TOKEN] - remove_duplicate = self.config['remove_duplicates_in_prediction'] + remove_duplicate = self.config["remove_duplicates_in_prediction"] variable_nums = [] for ast_id, example in enumerate(examples): variable_nums.append(len(example.ast.variables)) - beams = OrderedDict((ast_id, [self.Hypothesis([[]], 0, 0.)]) for ast_id in range(batch_size)) + beams = OrderedDict( + (ast_id, [self.Hypothesis([[]], 0, 0.0)]) for ast_id in range(batch_size) + ) hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) completed_hyps = [[] for _ in range(batch_size)] tgt_vocab_size = len(self.vocab.target) @@ -99,24 +125,31 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: attention_memory = self.get_attention_memory(context_encoding) context_encoding_t = attention_memory - h_tm1 = h_0 = self.get_init_state(context_encoding) + h_tm1 = _h_0 = self.get_init_state(context_encoding) - # Note that we are using the `restoration_indices` from `context_encoding`, which is the word-level restoration index + # Note that we are using the `restoration_indices` from + # `context_encoding`, which is the word-level restoration index # (batch_size, variable_master_node_num, encoding_size) - variable_encoding = context_encoding['variable_encoding'] + variable_encoding = context_encoding["variable_encoding"] # (batch_size, encoding_size) - variable_name_embed_tm1 = att_tm1 = torch.zeros(batch_size, self.lstm_cell.hidden_size, device=self.device) + variable_name_embed_tm1 = att_tm1 = torch.zeros( + batch_size, self.lstm_cell.hidden_size, device=self.device + ) - max_prediction_time_step = self.config['max_prediction_time_step'] + max_prediction_time_step = self.config["max_prediction_time_step"] for t in range(0, max_prediction_time_step): # (total_live_hyp_num, encoding_size) if t > 0: - variable_encoding_t = variable_encoding[hyp_ast_ids_t, hyp_variable_ptrs_t] + variable_encoding_t = variable_encoding[ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] else: variable_encoding_t = variable_encoding[:, 0] - if self.config['input_feed']: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1) + if self.config["input_feed"]: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 + ) else: x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) @@ -139,12 +172,15 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: for beam_id, (ast_id, beam) in enumerate(beams.items()): beam_end_hyp_pos = beam_start_hyp_pos + len(beam) # (live_beam_size, vocab_size) - beam_cont_cand_hyp_scores = cont_cand_hyp_scores[beam_start_hyp_pos: beam_end_hyp_pos] + beam_cont_cand_hyp_scores = cont_cand_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] cont_beam_size = beam_size - len(completed_hyps[ast_id]) # Take `len(beam)` more candidates to account for possible duplicate k = min(beam_cont_cand_hyp_scores.numel(), cont_beam_size + len(beam)) beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( - beam_cont_cand_hyp_scores.view(-1), k=k, dim=-1) + beam_cont_cand_hyp_scores.view(-1), k=k, dim=-1 + ) # (cont_beam_size) beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size @@ -163,7 +199,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: variable_ptr = prev_hyp.variable_ptr new_variable_list = list(prev_hyp.variable_list) - new_variable_list[-1] = list(new_variable_list[-1] + [hyp_var_name_id]) + new_variable_list[-1] = list( + new_variable_list[-1] + [hyp_var_name_id] + ) if hyp_var_name_id == end_of_variable_id: # remove empty cases @@ -172,17 +210,22 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if remove_duplicate: last_pred = new_variable_list[-1] - if any(x == last_pred for x in new_variable_list[:-1] if x != [same_variable_id, end_of_variable_id]): - # print('found a duplicate!', ', '.join([str(x) for x in last_pred])) + if any( + x == last_pred + for x in new_variable_list[:-1] + if x != [same_variable_id, end_of_variable_id] + ): continue variable_ptr += 1 new_variable_list.append([]) beam_cnt += 1 - new_hyp = self.Hypothesis(variable_list=new_variable_list, - variable_ptr=variable_ptr, - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=new_variable_list, + variable_ptr=variable_ptr, + score=new_hyp_score, + ) if variable_ptr == variable_nums[ast_id]: completed_hyps[ast_id].append(new_hyp) @@ -194,7 +237,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: new_hyp_var_name_ids.append(hyp_var_name_id) new_hyp_ast_ids.append(ast_id) new_hyp_variable_ptrs.append(variable_ptr) - is_same_variable_mask.append(1. if prev_hyp.variable_ptr == variable_ptr else 0.) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) if beam_cnt >= cont_beam_size: break @@ -206,10 +251,14 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: h_tm1 = (h_t[0][live_prev_hyp_ids], h_t[1][live_prev_hyp_ids]) att_tm1 = q_t[live_prev_hyp_ids] - if self.config['tie_embedding']: - variable_name_embed_tm1 = self.state2names.weight[new_hyp_var_name_ids] + if self.config["tie_embedding"]: + variable_name_embed_tm1 = self.state2names.weight[ + new_hyp_var_name_ids + ] else: - variable_name_embed_tm1 = self.var_name_embed.weight[new_hyp_var_name_ids] + variable_name_embed_tm1 = self.var_name_embed.weight[ + new_hyp_var_name_ids + ] hyp_ast_ids_t = new_hyp_ast_ids hyp_variable_ptrs_t = new_hyp_variable_ptrs @@ -217,15 +266,26 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: beams = new_beams # (total_hyp_num, max_tree_size, node_encoding_size) - context_encoding_t = dict(attention_key=attention_memory['attention_key'][hyp_ast_ids_t], - attention_value=attention_memory['attention_value'][hyp_ast_ids_t], - attention_value_mask=attention_memory['attention_value_mask'][hyp_ast_ids_t]) + context_encoding_t = dict( + attention_key=attention_memory["attention_key"][hyp_ast_ids_t], + attention_value=attention_memory["attention_value"][hyp_ast_ids_t], + attention_value_mask=attention_memory["attention_value_mask"][ + hyp_ast_ids_t + ], + ) if self.independent_prediction_for_each_variable: - is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) - h_tm1 = (h_tm1[0] * is_same_variable_mask, h_tm1[1] * is_same_variable_mask) + is_same_variable_mask = torch.tensor( + is_same_variable_mask, device=self.device, dtype=torch.float + ).unsqueeze(-1) + h_tm1 = ( + h_tm1[0] * is_same_variable_mask, + h_tm1[1] * is_same_variable_mask, + ) att_tm1 = att_tm1 * is_same_variable_mask - variable_name_embed_tm1 = variable_name_embed_tm1 * is_same_variable_mask + variable_name_embed_tm1 = ( + variable_name_embed_tm1 * is_same_variable_mask + ) else: break @@ -236,32 +296,19 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if not hyps: # return identity renamings - print(f'Failed to found a hypothesis for function {ast.compilation_unit}', file=sys.stderr) + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) variable_rename_result = dict() for old_name in ast.variables: - variable_rename_result[old_name] = {'new_name': old_name, - 'prob': 0.} + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } example_rename_results = [variable_rename_result] else: - # top_hyp = hyps[0] - # sub_token_ptr = 0 - # for old_name in ast.variables: - # sub_token_begin = sub_token_ptr - # while top_hyp.variable_list[sub_token_ptr] != end_of_variable_id: - # sub_token_ptr += 1 - # sub_token_ptr += 1 # point to first sub-token of next variable - # sub_token_end = sub_token_ptr - # - # var_name_token_ids = top_hyp.variable_list[sub_token_begin: sub_token_end] # include ending - # if var_name_token_ids == [same_variable_id, end_of_variable_id]: - # new_var_name = old_name - # else: - # new_var_name = self.vocab.target.subtoken_model.decode_ids(var_name_token_ids) - # - # variable_rename_result[old_name] = {'new_name': new_var_name, - # 'prob': top_hyp.score} - example_rename_results = [] for hyp in hyps: @@ -271,9 +318,14 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if var_name_token_ids == [same_variable_id, end_of_variable_id]: new_var_name = old_name else: - new_var_name = self.vocab.target.subtoken_model.decode_ids(var_name_token_ids) - - variable_rename_result[old_name] = {'new_name': new_var_name, 'prob': hyp.score} + new_var_name = self.vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": hyp.score, + } example_rename_results.append(variable_rename_result) diff --git a/neural-model/model/embedding.py b/neural-model/model/embedding.py index 7f38a47..3577d69 100644 --- a/neural-model/model/embedding.py +++ b/neural-model/model/embedding.py @@ -1,10 +1,8 @@ -from typing import List, Dict +from typing import Dict, List import sentencepiece as spm - import torch import torch.nn as nn - from utils.nn_util import SMALL_NUMBER @@ -16,16 +14,25 @@ def __init__(self, bpe_model_path, embedding_size): self.bpe_model.load(bpe_model_path) self.size = len(self.bpe_model) self.pad_id = self.bpe_model.pad_id() - self.embeddings = nn.Embedding(self.size, embedding_size, padding_idx=self.pad_id) + self.embeddings = nn.Embedding( + self.size, embedding_size, padding_idx=self.pad_id + ) @property def device(self): return self.embeddings.weight.device @classmethod - def to_input_tensor(cls, sub_tokens_list: List[List[str]], bpe_model: spm.SentencePieceProcessor, pad_id=0) -> torch.Tensor: + def to_input_tensor( + cls, + sub_tokens_list: List[List[str]], + bpe_model: spm.SentencePieceProcessor, + pad_id=0, + ) -> torch.Tensor: max_subword_num = max(len(x) for x in sub_tokens_list) - idx_tensor = torch.zeros(len(sub_tokens_list), max_subword_num, dtype=torch.long) + idx_tensor = torch.zeros( + len(sub_tokens_list), max_subword_num, dtype=torch.long + ) idx_tensor.fill_(pad_id) for i, token_list in enumerate(sub_tokens_list): @@ -55,22 +62,28 @@ def __init__(self, type_num, embedding_size, pad_id=0): self.size = type_num self.pad_id = pad_id - self.embeddings = nn.Embedding(type_num, embedding_size, padding_idx=self.pad_id) + self.embeddings = nn.Embedding( + type_num, embedding_size, padding_idx=self.pad_id + ) @property def device(self): return self.embeddings.weight.device @classmethod - def to_input_tensor(cls, token_types_list: List[List[str]], - type2id: Dict[str, int], - pad_id=0) -> torch.Tensor: + def to_input_tensor( + cls, token_types_list: List[List[str]], type2id: Dict[str, int], pad_id=0 + ) -> torch.Tensor: max_subword_num = max(len(x) for x in token_types_list) - idx_tensor = torch.zeros(len(token_types_list), max_subword_num, dtype=torch.long) + idx_tensor = torch.zeros( + len(token_types_list), max_subword_num, dtype=torch.long + ) idx_tensor.fill_(pad_id) for i, token_list in enumerate(token_types_list): - idx_tensor[i, :len(token_list)] = torch.tensor([type2id[t] for t in token_list]) + idx_tensor[i, : len(token_list)] = torch.tensor( + [type2id[t] for t in token_list] + ) return idx_tensor @@ -78,6 +91,8 @@ def forward(self, type_tokens_indices: torch.Tensor) -> torch.Tensor: # type_tokens_indices: (batch_size, max_sub_token_num) sub_tokens_mask = torch.ne(type_tokens_indices, self.pad_id).float() embedding = self.embeddings(type_tokens_indices) * sub_tokens_mask.unsqueeze(-1) - embedding = embedding.sum(dim=1) / (sub_tokens_mask.sum(-1) + SMALL_NUMBER).unsqueeze(-1) + embedding = embedding.sum(dim=1) / ( + sub_tokens_mask.sum(-1) + SMALL_NUMBER + ).unsqueeze(-1) return embedding diff --git a/neural-model/model/encoder.py b/neural-model/model/encoder.py index f73cbf2..7ab869e 100644 --- a/neural-model/model/encoder.py +++ b/neural-model/model/encoder.py @@ -3,4 +3,3 @@ class Encoder(nn.Module): pass - diff --git a/neural-model/model/ensemble_model.py b/neural-model/model/ensemble_model.py index 8e2c054..c915adc 100644 --- a/neural-model/model/ensemble_model.py +++ b/neural-model/model/ensemble_model.py @@ -13,29 +13,28 @@ def __init__(self, model_paths, use_cuda, config): model = RenamingModel.load(model_path, use_cuda=use_cuda) self.models.append(model) - self.max_prediction_time_step = min(model.decoder.config['max_prediction_time_step'] for model in self.models) - self.device = torch.device('cuda:0') if use_cuda else torch.device('cpu') - - self.Hypothesis = namedtuple('Hypothesis', ['variable_list', 'variable_ptr', 'score']) + self.max_prediction_time_step = min( + model.decoder.config["max_prediction_time_step"] for model in self.models + ) + self.device = torch.device("cuda:0") if use_cuda else torch.device("cpu") + + self.Hypothesis = namedtuple( + "Hypothesis", ["variable_list", "variable_ptr", "score"] + ) self.config = config - print('Ensemble Model Config', file=sys.stderr) + print("Ensemble Model Config", file=sys.stderr) print(self.config, file=sys.stderr) @classmethod def default_params(cls): return { - 'encoder': { - "type": "EnsembleModel" - }, - 'decoder': { - 'beam_size': 5, - 'remove_duplicates_in_prediction': False - }, - 'train': { - 'eval_batch_size': 100, - 'num_batchers': 5, - 'num_readers': 5, - 'buffer_size': 100 + "encoder": {"type": "EnsembleModel"}, + "decoder": {"beam_size": 5, "remove_duplicates_in_prediction": False}, + "train": { + "eval_batch_size": 100, + "num_batchers": 5, + "num_readers": 5, + "buffer_size": 100, }, } @@ -50,7 +49,7 @@ def predict(self, examples): context_encodings = [] attention_memories = [] variable_encodings = [] - remove_duplicate = self.config['decoder']['remove_duplicates_in_prediction'] + remove_duplicate = self.config["decoder"]["remove_duplicates_in_prediction"] h_tm1 = [] variable_name_embed_tm1 = [] @@ -64,23 +63,27 @@ def predict(self, examples): h_0 = model.decoder.get_init_state(context_encoding) context_encodings.append(context_encoding) - variable_encodings.append(context_encoding['variable_encoding']) + variable_encodings.append(context_encoding["variable_encoding"]) attention_memories.append(attention_memory) h_tm1.append(h_0) - att_tm1 = torch.zeros(len(examples), model.decoder.lstm_cell.hidden_size, device=self.device) + att_tm1 = torch.zeros( + len(examples), model.decoder.lstm_cell.hidden_size, device=self.device + ) variable_name_embed_tm1.append(att_tm1) context_encoding_t = attention_memories - beam_size = self.config['decoder']['beam_size'] + beam_size = self.config["decoder"]["beam_size"] batch_size = len(examples) vocab = self.models[0].decoder.vocab same_variable_id = vocab.target[SAME_VARIABLE_TOKEN] end_of_variable_id = vocab.target[END_OF_VARIABLE_TOKEN] variable_nums = [len(e.ast.variables) for e in examples] - beams = OrderedDict((ast_id, [self.Hypothesis([[]], 0, 0.)]) for ast_id in range(batch_size)) + beams = OrderedDict( + (ast_id, [self.Hypothesis([[]], 0, 0.0)]) for ast_id in range(batch_size) + ) hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) completed_hyps = [[] for _ in range(batch_size)] tgt_vocab_size = len(self.models[0].decoder.vocab.target) @@ -92,26 +95,45 @@ def predict(self, examples): for model_id, model in enumerate(self.models): # (total_live_hyp_num, encoding_size) if t > 0: - variable_encoding_t = variable_encodings[model_id][hyp_ast_ids_t, hyp_variable_ptrs_t] + variable_encoding_t = variable_encodings[model_id][ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] else: variable_encoding_t = variable_encodings[model_id][:, 0] - if model.decoder.config['input_feed']: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1[model_id], att_tm1], dim=-1) + if model.decoder.config["input_feed"]: + x = torch.cat( + [ + variable_encoding_t, + variable_name_embed_tm1[model_id], + att_tm1, + ], + dim=-1, + ) else: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1[model_id]], dim=-1) + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1[model_id]], dim=-1 + ) - h_t, q_t, alpha_t = model.decoder.rnn_step(x, h_tm1[model_id], context_encoding_t[model_id]) + h_t, q_t, alpha_t = model.decoder.rnn_step( + x, h_tm1[model_id], context_encoding_t[model_id] + ) model_h_t.append(h_t) model_q_t.append(q_t) # (total_live_hyp_num, vocab_size) - hyp_var_name_scores_t = torch.log_softmax(model.decoder.state2names(q_t), dim=-1) - cont_cand_hyp_scores = hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t + hyp_var_name_scores_t = torch.log_softmax( + model.decoder.state2names(q_t), dim=-1 + ) + cont_cand_hyp_scores = ( + hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t + ) candidate_cont_hyp_scores.append(cont_cand_hyp_scores) # (total_live_hyp_num, vocab_size) - candidate_cont_hyp_scores = torch.logsumexp(torch.stack(candidate_cont_hyp_scores, dim=-1), dim=-1) + candidate_cont_hyp_scores = torch.logsumexp( + torch.stack(candidate_cont_hyp_scores, dim=-1), dim=-1 + ) new_beams = OrderedDict() live_beam_ids = [] @@ -126,12 +148,14 @@ def predict(self, examples): for beam_id, (ast_id, beam) in enumerate(beams.items()): beam_end_hyp_pos = beam_start_hyp_pos + len(beam) # (live_beam_size, vocab_size) - beam_cont_cand_hyp_scores = candidate_cont_hyp_scores[beam_start_hyp_pos: beam_end_hyp_pos] + beam_cont_cand_hyp_scores = candidate_cont_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] cont_beam_size = beam_size - len(completed_hyps[ast_id]) - beam_new_hyp_scores, beam_new_hyp_positions = torch.topk(beam_cont_cand_hyp_scores.view(-1), - k=cont_beam_size, - dim=-1) + beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( + beam_cont_cand_hyp_scores.view(-1), k=cont_beam_size, dim=-1 + ) # (cont_beam_size) beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size @@ -149,7 +173,9 @@ def predict(self, examples): variable_ptr = prev_hyp.variable_ptr new_variable_list = list(prev_hyp.variable_list) - new_variable_list[-1] = list(new_variable_list[-1] + [hyp_var_name_id]) + new_variable_list[-1] = list( + new_variable_list[-1] + [hyp_var_name_id] + ) if hyp_var_name_id == end_of_variable_id: # remove empty cases @@ -158,16 +184,22 @@ def predict(self, examples): if remove_duplicate: last_pred = new_variable_list[-1] - if any(x == last_pred for x in new_variable_list[:-1] if x != [same_variable_id, end_of_variable_id]): + if any( + x == last_pred + for x in new_variable_list[:-1] + if x != [same_variable_id, end_of_variable_id] + ): # print('found a duplicate!', ', '.join([str(x) for x in last_pred])) continue variable_ptr += 1 new_variable_list.append([]) - new_hyp = self.Hypothesis(variable_list=new_variable_list, - variable_ptr=variable_ptr, - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=new_variable_list, + variable_ptr=variable_ptr, + score=new_hyp_score, + ) if variable_ptr == variable_nums[ast_id]: completed_hyps[ast_id].append(new_hyp) @@ -179,29 +211,50 @@ def predict(self, examples): new_hyp_var_name_ids.append(hyp_var_name_id) new_hyp_ast_ids.append(ast_id) new_hyp_variable_ptrs.append(variable_ptr) - is_same_variable_mask.append(1. if prev_hyp.variable_ptr == variable_ptr else 0.) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) beam_start_hyp_pos = beam_end_hyp_pos if live_beam_ids: hyp_scores_tm1 = torch.tensor(new_hyp_scores, device=self.device) - h_tm1 = [(model_h_t[model_id][0][live_prev_hyp_ids], model_h_t[model_id][1][live_prev_hyp_ids]) - for model_id in range(len(self.models))] - att_tm1 = [model_q_t[model_id][live_prev_hyp_ids] - for model_id in range(len(self.models))] - - variable_name_embed_tm1 = [model.decoder.state2names.weight[new_hyp_var_name_ids] - for model in self.models] + h_tm1 = [ + ( + model_h_t[model_id][0][live_prev_hyp_ids], + model_h_t[model_id][1][live_prev_hyp_ids], + ) + for model_id in range(len(self.models)) + ] + att_tm1 = [ + model_q_t[model_id][live_prev_hyp_ids] + for model_id in range(len(self.models)) + ] + + variable_name_embed_tm1 = [ + model.decoder.state2names.weight[new_hyp_var_name_ids] + for model in self.models + ] hyp_ast_ids_t = new_hyp_ast_ids hyp_variable_ptrs_t = new_hyp_variable_ptrs beams = new_beams # (total_hyp_num, max_tree_size, node_encoding_size) - context_encoding_t = [dict(attention_key=attention_memories[model_id]['attention_key'][hyp_ast_ids_t], - attention_value=attention_memories[model_id]['attention_value'][hyp_ast_ids_t], - attention_value_mask=attention_memories[model_id]['attention_value_mask'][hyp_ast_ids_t]) - for model_id in range(len(self.models))] + context_encoding_t = [ + dict( + attention_key=attention_memories[model_id]["attention_key"][ + hyp_ast_ids_t + ], + attention_value=attention_memories[model_id]["attention_value"][ + hyp_ast_ids_t + ], + attention_value_mask=attention_memories[model_id][ + "attention_value_mask" + ][hyp_ast_ids_t], + ) + for model_id in range(len(self.models)) + ] # if self.independent_prediction_for_each_variable: # is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) @@ -219,10 +272,15 @@ def predict(self, examples): if not hyps: # return identity renamings - print(f'Failed to found a hypothesis for function {ast.compilation_unit}', file=sys.stderr) + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) for old_name in ast.variables: - variable_rename_result[old_name] = {'new_name': old_name, - 'prob': 0.} + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } else: top_hyp = hyps[0] # sub_token_ptr = 0 @@ -247,9 +305,14 @@ def predict(self, examples): if var_name_token_ids == [same_variable_id, end_of_variable_id]: new_var_name = old_name else: - new_var_name = vocab.target.subtoken_model.decode_ids(var_name_token_ids) - - variable_rename_result[old_name] = {'new_name': new_var_name, 'prob': top_hyp.score} + new_var_name = vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } variable_rename_results.append(variable_rename_result) diff --git a/neural-model/model/gnn.py b/neural-model/model/gnn.py index 893c817..8f6766c 100644 --- a/neural-model/model/gnn.py +++ b/neural-model/model/gnn.py @@ -6,10 +6,11 @@ class AdjacencyList: """represent the topology of a graph""" + def __init__(self, node_num: int, adj_list: List, device: torch.device = None): self.node_num = node_num if device is None: - device = torch.device('cpu') + device = torch.device("cpu") self.data = torch.tensor(adj_list, dtype=torch.long, device=device) self.edge_num = len(adj_list) @@ -17,7 +18,7 @@ def __init__(self, node_num: int, adj_list: List, device: torch.device = None): def device(self): return self.data.device - def to(self, device: torch.device) -> 'AdjacencyList': + def to(self, device: torch.device) -> "AdjacencyList": if self.data.device != device: self.data = self.data.to(device) return self @@ -27,11 +28,16 @@ def __getitem__(self, item): class GatedGraphNeuralNetwork(nn.Module): - def __init__(self, hidden_size, num_edge_types, layer_timesteps, - residual_connections, - state_to_message_dropout=0.3, - rnn_dropout=0.3, - use_bias_for_message_linear=True): + def __init__( + self, + hidden_size, + num_edge_types, + layer_timesteps, + residual_connections, + state_to_message_dropout=0.3, + rnn_dropout=0.3, + use_bias_for_message_linear=True, + ): super(GatedGraphNeuralNetwork, self).__init__() self.hidden_size = hidden_size @@ -42,19 +48,29 @@ def __init__(self, hidden_size, num_edge_types, layer_timesteps, self.rnn_dropout = rnn_dropout self.use_bias_for_message_linear = use_bias_for_message_linear - # Prepare linear transformations from node states to messages, for each layer and each edge type - # Prepare rnn cells for each layer + # Prepare linear transformations from node states to messages, for each + # layer and each edge type. + # Prepare rnn cells for each layer. self.state_to_message_linears = nn.ModuleDict() self.rnn_cells = nn.ModuleList() for layer_idx in range(len(self.layer_timesteps)): # Initiate a linear transformation for each edge type for edge_type_j in range(self.num_edge_types): - state_to_msg_linear_layer_i_type_j = nn.Linear(self.hidden_size, self.hidden_size, bias=use_bias_for_message_linear) - layer_name = 'state_to_message_linear_layer%d_type%d' % (layer_idx, edge_type_j) - self.state_to_message_linears[layer_name] = state_to_msg_linear_layer_i_type_j + state_to_msg_linear_layer_i_type_j = nn.Linear( + self.hidden_size, self.hidden_size, bias=use_bias_for_message_linear + ) + layer_name = ( + f"state_to_message_linear_layer{layer_idx}" f"_type{edge_type_j}" + ) + self.state_to_message_linears[ + layer_name + ] = state_to_msg_linear_layer_i_type_j layer_residual_connections = self.residual_connections.get(layer_idx, []) - rnn_cell_layer_i = nn.GRUCell(self.hidden_size * (1 + len(layer_residual_connections)), self.hidden_size) + rnn_cell_layer_i = nn.GRUCell( + self.hidden_size * (1 + len(layer_residual_connections)), + self.hidden_size, + ) self.rnn_cells.append(rnn_cell_layer_i) self.state_to_message_dropout_layer = nn.Dropout(self.state_to_message_dropout) @@ -64,43 +80,54 @@ def __init__(self, hidden_size, num_edge_types, layer_timesteps, def device(self): return self.rnn_cells[0].weight_hh.device - def forward(self, - initial_node_representation: torch.Tensor, - adjacency_lists: List[AdjacencyList], - return_all_states=False) -> torch.Tensor: - return self.compute_node_representations(initial_node_representation, adjacency_lists, - return_all_states=return_all_states) - - def compute_node_representations(self, - initial_node_representation: torch.Tensor, - adjacency_lists: List[AdjacencyList], - return_all_states=False) -> torch.Tensor: + def forward( + self, + initial_node_representation: torch.Tensor, + adjacency_lists: List[AdjacencyList], + return_all_states=False, + ) -> torch.Tensor: + return self.compute_node_representations( + initial_node_representation, + adjacency_lists, + return_all_states=return_all_states, + ) + + def compute_node_representations( + self, + initial_node_representation: torch.Tensor, + adjacency_lists: List[AdjacencyList], + return_all_states=False, + ) -> torch.Tensor: """ V: number of nodes on the batched graph E: number of edges on the batched graph D: hidden size of RNNs and output M: total number of messages """ - # If the dimension of initial node embedding is smaller, then perform padding first + # If the dimension of initial node embedding is smaller, then perform + # padding first node_num = initial_node_representation.size(0) init_node_repr_size = initial_node_representation.size(1) if init_node_repr_size < self.hidden_size: pad_size = self.hidden_size - init_node_repr_size - zero_pads = torch.zeros(node_num, pad_size, dtype=torch.float, device=self.device) - initial_node_representation = torch.cat([initial_node_representation, zero_pads], dim=-1) + zero_pads = torch.zeros( + node_num, pad_size, dtype=torch.float, device=self.device + ) + initial_node_representation = torch.cat( + [initial_node_representation, zero_pads], dim=-1 + ) # Shape: [V, D] node_states_per_layer = [initial_node_representation] - message_targets = [] # list of tensors of message targets of shape [E] + # list of tensors of message targets of shape [E] + message_targets = [] for edge_type_idx, adjacency_list_for_edge_type in enumerate(adjacency_lists): if adjacency_list_for_edge_type.edge_num > 0: edge_targets = adjacency_list_for_edge_type[:, 1] message_targets.append(edge_targets) message_targets = torch.cat(message_targets, dim=0) # Shape [M] - # sparse matrix of shape [V, M] - # incoming_msg_sparse_matrix = self.get_incoming_message_sparse_matrix(adjacency_lists).to(device) for layer_idx, num_timesteps in enumerate(self.layer_timesteps): # Used shape abbreviations: # V ~ number of nodes @@ -111,29 +138,41 @@ def compute_node_representations(self, # Extract residual messages, if any: layer_residual_connections = self.residual_connections.get(layer_idx, []) # List[(V, D)] - layer_residual_states: List[torch.Tensor] = [node_states_per_layer[residual_layer_idx] - for residual_layer_idx in layer_residual_connections] + layer_residual_states: List[torch.Tensor] = [ + node_states_per_layer[residual_layer_idx] + for residual_layer_idx in layer_residual_connections + ] - # Record new states for this layer. Initialised to last state, but will be updated below: + # Record new states for this layer. Initialised to last state, but + # will be updated below: node_states_for_this_layer = node_states_per_layer[-1] # For each message propagation step for t in range(num_timesteps): - messages: List[torch.Tensor] = [] # list of tensors of messages of shape [E, D] - message_source_states: List[torch.Tensor] = [] # list of tensors of edge source states of shape [E, D] + # list of tensors of messages of shape [E, D] + messages: List[torch.Tensor] = [] + # list of tensors of edge source states of shape [E, D] + message_source_states: List[torch.Tensor] = [] # Collect incoming messages per edge type - for edge_type_idx, adjacency_list_for_edge_type in enumerate(adjacency_lists): + for edge_type_idx, adjacency_list_for_edge_type in enumerate( + adjacency_lists + ): if adjacency_list_for_edge_type.edge_num > 0: # shape [E] edge_sources = adjacency_list_for_edge_type[:, 0] # shape [E, D] edge_source_states = node_states_for_this_layer[edge_sources] - layer_name = 'state_to_message_linear_layer%d_type%d' % (layer_idx, edge_type_idx) + layer_name = ( + f"state_to_message_linear_layer{layer_idx}" + f"_type{edge_type_idx}" + ) f_state_to_message = self.state_to_message_linears[layer_name] # Shape [E, D] - all_messages_for_edge_type = self.state_to_message_dropout_layer(f_state_to_message(edge_source_states)) + all_messages_for_edge_type = self.state_to_message_dropout_layer( + f_state_to_message(edge_source_states) + ) messages.append(all_messages_for_edge_type) message_source_states.append(edge_source_states) @@ -143,18 +182,28 @@ def compute_node_representations(self, # Sum up messages that go to the same target node # shape [V, D] - incoming_messages = torch.zeros(node_num, messages.size(1), device=self.device) - incoming_messages = incoming_messages.scatter_add_(0, - message_targets.unsqueeze(-1).expand_as(messages), - messages) - node_target_num = torch.zeros(node_num, device=self.device).scatter_add(0, message_targets, torch.ones(messages.size(0), device=self.device)) - incoming_messages = incoming_messages / (node_target_num + 1e-7).unsqueeze(-1) + incoming_messages = torch.zeros( + node_num, messages.size(1), device=self.device + ) + incoming_messages = incoming_messages.scatter_add_( + 0, message_targets.unsqueeze(-1).expand_as(messages), messages + ) + node_target_num = torch.zeros(node_num, device=self.device).scatter_add( + 0, message_targets, torch.ones(messages.size(0), device=self.device) + ) + incoming_messages = incoming_messages / ( + node_target_num + 1e-7 + ).unsqueeze(-1) # shape [V, D * (1 + num of residual connections)] - incoming_information = torch.cat(layer_residual_states + [incoming_messages], dim=-1) + incoming_information = torch.cat( + layer_residual_states + [incoming_messages], dim=-1 + ) # pass updated vertex features into RNN cell # Shape [V, D] - updated_node_states = self.rnn_cells[layer_idx](incoming_information, node_states_for_this_layer) + updated_node_states = self.rnn_cells[layer_idx]( + incoming_information, node_states_for_this_layer + ) updated_node_states = self.rnn_dropout_layer(updated_node_states) node_states_for_this_layer = updated_node_states @@ -168,17 +217,27 @@ def compute_node_representations(self, def main(): - gnn = GatedGraphNeuralNetwork(hidden_size=64, num_edge_types=2, - layer_timesteps=[3, 5, 7, 2], residual_connections={2: [0], 3: [0, 1]}) - - adj_list_type1 = AdjacencyList(node_num=4, adj_list=[(0, 2), (2, 1), (1, 3)], device=gnn.device) - adj_list_type2 = AdjacencyList(node_num=4, adj_list=[(0, 0), (0, 1)], device=gnn.device) - - node_representations = gnn.compute_node_representations(initial_node_representation=torch.randn(4, 64), - adjacency_lists=[adj_list_type1, adj_list_type2]) + gnn = GatedGraphNeuralNetwork( + hidden_size=64, + num_edge_types=2, + layer_timesteps=[3, 5, 7, 2], + residual_connections={2: [0], 3: [0, 1]}, + ) + + adj_list_type1 = AdjacencyList( + node_num=4, adj_list=[(0, 2), (2, 1), (1, 3)], device=gnn.device + ) + adj_list_type2 = AdjacencyList( + node_num=4, adj_list=[(0, 0), (0, 1)], device=gnn.device + ) + + node_representations = gnn.compute_node_representations( + initial_node_representation=torch.randn(4, 64), + adjacency_lists=[adj_list_type1, adj_list_type2], + ) print(node_representations) -if __name__ == '__main__': - main() \ No newline at end of file +if __name__ == "__main__": + main() diff --git a/neural-model/model/graph_encoder.py b/neural-model/model/graph_encoder.py index 9c1bbd0..18d758b 100644 --- a/neural-model/model/graph_encoder.py +++ b/neural-model/model/graph_encoder.py @@ -1,13 +1,13 @@ -from typing import List, Dict, Tuple +from operator import itemgetter +from typing import Dict, List, Tuple import numpy as np import torch -from torch import nn as nn - from model.embedding import NodeTypeEmbedder, SubTokenEmbedder from model.encoder import Encoder -from model.gnn import GatedGraphNeuralNetwork, AdjacencyList +from model.gnn import AdjacencyList, GatedGraphNeuralNetwork from model.sequential_encoder import SequentialEncoder +from torch import nn as nn from utils import util from utils.ast import AbstractSyntaxTree from utils.grammar import Grammar @@ -17,15 +17,18 @@ class GraphASTEncoder(Encoder): """An encoder based on gated recurrent graph neural networks""" - def __init__(self, - gnn: GatedGraphNeuralNetwork, - connections: List[str], - node_syntax_type_embedding_size: int, - decoder_hidden_size: int, - node_type_embedder: NodeTypeEmbedder, - node_content_embedder: SubTokenEmbedder, - vocab: Vocab, - config): + + def __init__( + self, + gnn: GatedGraphNeuralNetwork, + connections: List[str], + node_syntax_type_embedding_size: int, + decoder_hidden_size: int, + node_type_embedder: NodeTypeEmbedder, + node_content_embedder: SubTokenEmbedder, + vocab: Vocab, + config, + ): super(GraphASTEncoder, self).__init__() self.connections = connections @@ -34,28 +37,36 @@ def __init__(self, self.vocab = vocab self.grammar = grammar = vocab.grammar - self.node_syntax_type_embedding = nn.Embedding(len(grammar.syntax_types) + 2, node_syntax_type_embedding_size) + self.node_syntax_type_embedding = nn.Embedding( + len(grammar.syntax_types) + 2, node_syntax_type_embedding_size + ) self.variable_master_node_type_idx = len(grammar.syntax_types) self.master_node_type_idx = self.variable_master_node_type_idx + 1 - self.var_node_name_embedding = nn.Embedding(len(vocab.source), gnn.hidden_size, padding_idx=0) + self.var_node_name_embedding = nn.Embedding( + len(vocab.source), gnn.hidden_size, padding_idx=0 + ) self.node_type_embedder = node_type_embedder self.node_content_embedder = node_content_embedder - self.type_and_content_hybrid = nn.Linear(node_syntax_type_embedding_size + - node_type_embedder.embeddings.embedding_dim + - node_content_embedder.embeddings.embedding_dim, - gnn.hidden_size, bias=False) + self.type_and_content_hybrid = nn.Linear( + node_syntax_type_embedding_size + + node_type_embedder.embeddings.embedding_dim + + node_content_embedder.embeddings.embedding_dim, + gnn.hidden_size, + bias=False, + ) self.decoder_cell_init = nn.Linear(gnn.hidden_size, decoder_hidden_size) - self.init_with_seq_encoding = config['init_with_seq_encoding'] + self.init_with_seq_encoding = config["init_with_seq_encoding"] if self.init_with_seq_encoding: - self.seq_encoder = SequentialEncoder.build(config['seq_encoder']) - if config['seq_encoder']['source_encoding_size'] != gnn.hidden_size: - self.seq_variable_encoding_to_graph_linear = nn.Linear(config['seq_encoder']['source_encoding_size'], - gnn.hidden_size) + self.seq_encoder = SequentialEncoder.build(config["seq_encoder"]) + if config["seq_encoder"]["source_encoding_size"] != gnn.hidden_size: + self.seq_variable_encoding_to_graph_linear = nn.Linear( + config["seq_encoder"]["source_encoding_size"], gnn.hidden_size + ) else: self.seq_variable_encoding_to_graph_linear = lambda x: x @@ -68,54 +79,70 @@ def device(self): @classmethod def default_params(cls): return { - 'gnn': { - 'hidden_size': 128, - 'layer_timesteps': [8], - 'residual_connections': {'0': [0]} + "gnn": { + "hidden_size": 128, + "layer_timesteps": [8], + "residual_connections": {"0": [0]}, + }, + "connections": { + "top_down", + "bottom_up", + "variable_master_nodes", + "terminals", + "master_node", + "func_root_to_arg", }, - 'connections': {'top_down', 'bottom_up', 'variable_master_nodes', 'terminals', 'master_node', 'func_root_to_arg'}, - 'vocab_file': None, - 'bpe_model_path': None, - 'node_syntax_type_embedding_size': 64, - 'node_type_embedding_size': 64, - 'node_content_embedding_size': 128, - 'init_with_seq_encoding': False + "vocab_file": None, + "bpe_model_path": None, + "node_syntax_type_embedding_size": 64, + "node_type_embedding_size": 64, + "node_content_embedding_size": 128, + "init_with_seq_encoding": False, } @classmethod - def build(cls, config): + def build(cls, config, logging=False): params = util.update(GraphASTEncoder.default_params(), config) - print(params) + if logging: + print(params) - connections = params['connections'] + connections = params["connections"] connection2edge_type = { - 'top_down': 1, - 'bottom_up': 1, - 'variable_master_nodes': 2, - 'terminals': 2, - 'master_node': 2, - 'var_usage': 2, - 'func_root_to_arg': 1 + "top_down": 1, + "bottom_up": 1, + "variable_master_nodes": 2, + "terminals": 2, + "master_node": 2, + "var_usage": 2, + "func_root_to_arg": 1, } num_edge_types = sum(connection2edge_type[key] for key in connections) - gnn = GatedGraphNeuralNetwork(hidden_size=params['gnn']['hidden_size'], - layer_timesteps=params['gnn']['layer_timesteps'], - residual_connections=params['gnn']['residual_connections'], - num_edge_types=num_edge_types) - - vocab = Vocab.load(params['vocab_file']) - node_type_embedder = NodeTypeEmbedder(len(vocab.grammar.variable_types), params['node_type_embedding_size']) - node_content_embedder = SubTokenEmbedder(vocab.obj_name.subtoken_model_path, params['node_content_embedding_size']) - - model = cls(gnn, - params['connections'], - params['node_syntax_type_embedding_size'], - params['decoder_hidden_size'], - node_type_embedder, - node_content_embedder, - vocab, - config=params) + gnn = GatedGraphNeuralNetwork( + hidden_size=params["gnn"]["hidden_size"], + layer_timesteps=params["gnn"]["layer_timesteps"], + residual_connections=params["gnn"]["residual_connections"], + num_edge_types=num_edge_types, + ) + + vocab = Vocab.load(params["vocab_file"]) + node_type_embedder = NodeTypeEmbedder( + len(vocab.grammar.variable_types), params["node_type_embedding_size"] + ) + node_content_embedder = SubTokenEmbedder( + vocab.obj_name.subtoken_model_path, params["node_content_embedding_size"] + ) + + model = cls( + gnn, + params["connections"], + params["node_syntax_type_embedding_size"], + params["decoder_hidden_size"], + node_type_embedder, + node_content_embedder, + vocab, + config=params, + ) return model @@ -124,29 +151,48 @@ def forward(self, tensor_dict: Dict[str, torch.Tensor]): if self.init_with_seq_encoding: # scatter sequential encoding results to variable positions - seq_encoding = self.seq_encoder(tensor_dict['seq_encoder_input']) - seq_var_encoding = seq_encoding['variable_encoding'] - tree_node_init_encoding[tensor_dict['variable_master_node_ids']] = self.seq_variable_encoding_to_graph_linear( - seq_var_encoding.view(-1, seq_var_encoding.size(-1))[tensor_dict['var_repr_flattened_positions']]) + seq_encoding = self.seq_encoder(tensor_dict["seq_encoder_input"]) + seq_var_encoding = seq_encoding["variable_encoding"] + tree_node_init_encoding[ + tensor_dict["variable_master_node_ids"] + ] = self.seq_variable_encoding_to_graph_linear( + seq_var_encoding.view(-1, seq_var_encoding.size(-1))[ + tensor_dict["var_repr_flattened_positions"] + ] + ) # (batch_size, node_encoding_size) - tree_node_encoding = self.gnn.compute_node_representations(initial_node_representation=tree_node_init_encoding, - adjacency_lists=tensor_dict['adj_lists']) + tree_node_encoding = self.gnn.compute_node_representations( + initial_node_representation=tree_node_init_encoding, + adjacency_lists=tensor_dict["adj_lists"], + ) - connections = self.config['connections'] + connections = self.config["connections"] - if 'variable_master_nodes' in connections: - variable_master_node_ids = tensor_dict['variable_master_node_ids'] + if "variable_master_nodes" in connections: + variable_master_node_ids = tensor_dict["variable_master_node_ids"] variable_encoding = tree_node_encoding[variable_master_node_ids] else: - var_nodes_encoding = tree_node_encoding[tensor_dict['variable_node_ids']] - variable_num = tensor_dict['variable_mention_nums'].size(0) - variable_encoding = torch.zeros(variable_num, var_nodes_encoding.size(-1), device=self.device) - var_nodes_encoding_sum = variable_encoding.scatter_add_(0, tensor_dict['variable_node_variable_ids'].unsqueeze(-1).expand_as(var_nodes_encoding), var_nodes_encoding) - variable_encoding = var_nodes_encoding_sum / tensor_dict['variable_mention_nums'].unsqueeze(-1) + var_nodes_encoding = tree_node_encoding[tensor_dict["variable_node_ids"]] + variable_num = tensor_dict["variable_mention_nums"].size(0) + variable_encoding = torch.zeros( + variable_num, var_nodes_encoding.size(-1), device=self.device + ) + var_nodes_encoding_sum = variable_encoding.scatter_add_( + 0, + tensor_dict["variable_node_variable_ids"] + .unsqueeze(-1) + .expand_as(var_nodes_encoding), + var_nodes_encoding, + ) + variable_encoding = var_nodes_encoding_sum / tensor_dict[ + "variable_mention_nums" + ].unsqueeze(-1) # restore variable encoding - variable_encoding = variable_encoding[tensor_dict['variable_encoding_restoration_indices']] * tensor_dict['variable_encoding_restoration_indices_mask'].unsqueeze(-1) + variable_encoding = variable_encoding[ + tensor_dict["variable_encoding_restoration_indices"] + ] * tensor_dict["variable_encoding_restoration_indices_mask"].unsqueeze(-1) context_encoding = dict( tree_node_init_encoding=tree_node_init_encoding, @@ -160,8 +206,12 @@ def forward(self, tensor_dict: Dict[str, torch.Tensor]): # noinspection PyUnboundLocalVariable @classmethod - def to_packed_graph(cls, asts: List[AbstractSyntaxTree], connections: List, - init_with_seq_encoding: bool = False) -> Tuple[PackedGraph, Dict]: + def to_packed_graph( + cls, + asts: List[AbstractSyntaxTree], + connections: List, + init_with_seq_encoding: bool = False, + ) -> Tuple[PackedGraph, Dict]: packed_graph = PackedGraph(asts) max_variable_num = max(len(ast.variables) for ast in asts) @@ -171,15 +221,22 @@ def to_packed_graph(cls, asts: List[AbstractSyntaxTree], connections: List, var_master_nodes_adj_list = [] var_usage_adj_list = [] func_root_to_arg_adj_list = [] - - var_node_ids = [] # list of node ids of variable mentions - var_repr_flattened_positions = [] # list of positions of variable encodings in the flattened version of restored `variable_encoding` + # list of node ids of variable mentions + var_node_ids = [] + # list of positions of variable encodings in the flattened version of + # restored `variable_encoding` + var_repr_flattened_positions = [] var_node_variable_ids = [] - var_mention_nums = [] # list of number of mentions for each variable + # list of number of mentions for each variable + var_mention_nums = [] - use_variable_master_node = 'variable_master_nodes' in connections - variable_repr_restoration_indices = torch.zeros(len(asts), max_variable_num, dtype=torch.long) - variable_repr_restoration_indices_mask = torch.zeros(len(asts), max_variable_num) + use_variable_master_node = "variable_master_nodes" in connections + variable_repr_restoration_indices = torch.zeros( + len(asts), max_variable_num, dtype=torch.long + ) + variable_repr_restoration_indices_mask = torch.zeros( + len(asts), max_variable_num + ) variable_cum_count = 0 for ast_id, ast in enumerate(asts): @@ -189,49 +246,78 @@ def to_packed_graph(cls, asts: List[AbstractSyntaxTree], connections: List, node_adj_list.append((prev_node_packed_id, succ_node_packed_id)) - if 'terminals' in connections: + if "terminals" in connections: for i in range(len(ast.terminal_nodes) - 1): terminal_i = ast.terminal_nodes[i] terminal_ip1 = ast.terminal_nodes[i + 1] - terminal_nodes_adj_list.append(( - packed_graph.get_packed_node_id(ast_id, terminal_i), - packed_graph.get_packed_node_id(ast_id, terminal_ip1) - )) + terminal_nodes_adj_list.append( + ( + packed_graph.get_packed_node_id(ast_id, terminal_i), + packed_graph.get_packed_node_id(ast_id, terminal_ip1), + ) + ) # add master node - if 'master_node' in connections: - master_node_id = packed_graph.register_node(ast_id, 'master_node', group='master_nodes') - - master_node_adj_list.extend([ - (master_node_id, packed_graph.get_packed_node_id(ast_id, node)) - for node in ast - ]) + if "master_node" in connections: + master_node_id = packed_graph.register_node( + ast_id, "master_node", group="master_nodes" + ) + + master_node_adj_list.extend( + [ + (master_node_id, packed_graph.get_packed_node_id(ast_id, node)) + for node in ast + ] + ) for i, (var_name, var_nodes) in enumerate(ast.variables.items()): var_node_num = len(var_nodes) # add data flow links - if 'var_usage' in connections: + if "var_usage" in connections: for _idx in range(len(var_nodes) - 1): _var_node, _next_var_node = var_nodes[_idx], var_nodes[_idx + 1] - var_usage_adj_list.append((packed_graph.get_packed_node_id(ast_id, _var_node), packed_graph.get_packed_node_id(ast_id, _next_var_node))) - - if 'func_root_to_arg' in connections and var_nodes[0].is_arg: - func_root_node_id = packed_graph.get_packed_node_id(ast_id, ast.root) + var_usage_adj_list.append( + ( + packed_graph.get_packed_node_id(ast_id, _var_node), + packed_graph.get_packed_node_id(ast_id, _next_var_node), + ) + ) + + if "func_root_to_arg" in connections and var_nodes[0].is_arg: + func_root_node_id = packed_graph.get_packed_node_id( + ast_id, ast.root + ) for var_node in var_nodes: - func_root_to_arg_adj_list.append((func_root_node_id, packed_graph.get_packed_node_id(ast_id, var_node))) + func_root_to_arg_adj_list.append( + ( + func_root_node_id, + packed_graph.get_packed_node_id(ast_id, var_node), + ) + ) if use_variable_master_node: - var_master_node_id, node_id_in_group = packed_graph.register_node(ast_id, var_name, - group='variable_master_nodes', - return_node_index_in_group=True) - - var_master_nodes_adj_list.extend([( - var_master_node_id, packed_graph.get_packed_node_id(ast_id, node)) - for node in var_nodes - ]) + var_master_node_id, node_id_in_group = packed_graph.register_node( + ast_id, + var_name, + group="variable_master_nodes", + return_node_index_in_group=True, + ) + + var_master_nodes_adj_list.extend( + [ + ( + var_master_node_id, + packed_graph.get_packed_node_id(ast_id, node), + ) + for node in var_nodes + ] + ) else: - var_node_packed_ids = [packed_graph.get_packed_node_id(ast_id, node) for node in var_nodes] + var_node_packed_ids = [ + packed_graph.get_packed_node_id(ast_id, node) + for node in var_nodes + ] var_node_ids.extend(var_node_packed_ids) var_node_variable_ids.extend([variable_cum_count] * var_node_num) var_mention_nums.append(var_node_num) @@ -241,83 +327,121 @@ def to_packed_graph(cls, asts: List[AbstractSyntaxTree], connections: List, variable_cum_count += 1 - variable_repr_restoration_indices_mask[ast_id, :len(ast.variables)] = 1. + variable_repr_restoration_indices_mask[ast_id, : len(ast.variables)] = 1.0 adj_lists = [] - if 'top_down' in connections: + if "top_down" in connections: adj_lists.append(node_adj_list) - if 'bottom_up' in connections: + if "bottom_up" in connections: reversed_node_adj_list = [(n2, n1) for n1, n2 in node_adj_list] adj_lists.append(reversed_node_adj_list) - if 'variable_master_nodes' in connections: - reversed_var_master_nodes_adj_list = [(n2, n1) for n1, n2 in var_master_nodes_adj_list] + if "variable_master_nodes" in connections: + reversed_var_master_nodes_adj_list = [ + (n2, n1) for n1, n2 in var_master_nodes_adj_list + ] adj_lists.append(master_node_adj_list) adj_lists.append(reversed_var_master_nodes_adj_list) - if 'var_usage' in connections: + if "var_usage" in connections: reversed_var_usage_adj_list = [(n2, n1) for n1, n2 in var_usage_adj_list] adj_lists.append(var_usage_adj_list) adj_lists.append(reversed_var_usage_adj_list) - if 'func_root_to_arg' in connections: + if "func_root_to_arg" in connections: adj_lists.append(func_root_to_arg_adj_list) - if 'terminals' in connections: - reversed_terminal_nodes_adj_list = [(n2, n1) for n1, n2 in terminal_nodes_adj_list] + if "terminals" in connections: + reversed_terminal_nodes_adj_list = [ + (n2, n1) for n1, n2 in terminal_nodes_adj_list + ] adj_lists.append(terminal_nodes_adj_list) adj_lists.append(reversed_terminal_nodes_adj_list) - if 'master_node' in connections: - reversed_master_node_adj_list = [(n2, n1) for n1, n2 in master_node_adj_list] + if "master_node" in connections: + reversed_master_node_adj_list = [ + (n2, n1) for n1, n2 in master_node_adj_list + ] adj_lists.append(master_node_adj_list) adj_lists.append(reversed_master_node_adj_list) - # if 'self_loop' in connections: - # self_loop_adj_list = [(n, n) for n in packed_graph.get_nodes_by_group('ast')] adj_lists = [ - AdjacencyList(adj_list=adj_list, node_num=packed_graph.size) for adj_list in adj_lists + AdjacencyList(adj_list=adj_list, node_num=packed_graph.size) + for adj_list in adj_lists ] max_tree_node_num = max(tree.size for tree in packed_graph.trees) - tree_restoration_indices = torch.zeros(packed_graph.tree_num, max_tree_node_num, dtype=torch.long) - tree_restoration_indices_mask = torch.zeros(packed_graph.tree_num, max_tree_node_num, dtype=torch.float) + tree_restoration_indices = torch.zeros( + packed_graph.tree_num, max_tree_node_num, dtype=torch.long + ) + tree_restoration_indices_mask = torch.zeros( + packed_graph.tree_num, max_tree_node_num, dtype=torch.float + ) - max_terminal_node_num = max(len(tree.terminal_nodes) for tree in packed_graph.trees) - terminal_node_restoration_indices = torch.zeros(packed_graph.tree_num, max_terminal_node_num, dtype=torch.long) - terminal_node_restoration_indices_mask = torch.zeros(packed_graph.tree_num, max_terminal_node_num, dtype=torch.float) + max_terminal_node_num = max( + len(tree.terminal_nodes) for tree in packed_graph.trees + ) + terminal_node_restoration_indices = torch.zeros( + packed_graph.tree_num, max_terminal_node_num, dtype=torch.long + ) + terminal_node_restoration_indices_mask = torch.zeros( + packed_graph.tree_num, max_terminal_node_num, dtype=torch.float + ) for ast_id in range(packed_graph.tree_num): - packed_node_ids = list(packed_graph.node_groups[ast_id]['ast_nodes'].values()) - tree_restoration_indices[ast_id, :len(packed_node_ids)] = torch.tensor(packed_node_ids) - tree_restoration_indices_mask[ast_id, :len(packed_node_ids)] = 1. + packed_node_ids = list( + packed_graph.node_groups[ast_id]["ast_nodes"].values() + ) + tree_restoration_indices[ast_id, : len(packed_node_ids)] = torch.tensor( + packed_node_ids + ) + tree_restoration_indices_mask[ast_id, : len(packed_node_ids)] = 1.0 tree = packed_graph.trees[ast_id] - terminal_node_ids = [packed_graph.get_packed_node_id(ast_id, node) for node in tree.terminal_nodes] - terminal_node_restoration_indices[ast_id, :len(terminal_node_ids)] = torch.tensor(terminal_node_ids) - terminal_node_restoration_indices_mask[ast_id, :len(terminal_node_ids)] = 1. - - tensor_dict = {'adj_lists': adj_lists, - 'variable_encoding_restoration_indices': variable_repr_restoration_indices, - 'variable_encoding_restoration_indices_mask': variable_repr_restoration_indices_mask, - 'tree_restoration_indices': tree_restoration_indices, - 'tree_restoration_indices_mask': tree_restoration_indices_mask, - 'terminal_node_restoration_indices': terminal_node_restoration_indices, - 'terminal_node_restoration_indices_mask': terminal_node_restoration_indices_mask, - 'packed_graph_size': packed_graph.size} + terminal_node_ids = [ + packed_graph.get_packed_node_id(ast_id, node) + for node in tree.terminal_nodes + ] + terminal_node_restoration_indices[ + ast_id, : len(terminal_node_ids) + ] = torch.tensor(terminal_node_ids) + terminal_node_restoration_indices_mask[ + ast_id, : len(terminal_node_ids) + ] = 1.0 + + tensor_dict = { + "adj_lists": adj_lists, + "variable_encoding_restoration_indices": variable_repr_restoration_indices, + "variable_encoding_restoration_indices_mask": variable_repr_restoration_indices_mask, + "tree_restoration_indices": tree_restoration_indices, + "tree_restoration_indices_mask": tree_restoration_indices_mask, + "terminal_node_restoration_indices": terminal_node_restoration_indices, + "terminal_node_restoration_indices_mask": terminal_node_restoration_indices_mask, + "packed_graph_size": packed_graph.size, + } if init_with_seq_encoding: - tensor_dict['var_repr_flattened_positions'] = torch.tensor(var_repr_flattened_positions, dtype=torch.long) + tensor_dict["var_repr_flattened_positions"] = torch.tensor( + var_repr_flattened_positions, dtype=torch.long + ) if use_variable_master_node: - tensor_dict['variable_master_node_ids'] = [_id for node, _id in - packed_graph.get_nodes_by_group('variable_master_nodes')] + tensor_dict["variable_master_node_ids"] = [ + _id + for node, _id in packed_graph.get_nodes_by_group( + "variable_master_nodes" + ) + ] else: - tensor_dict['variable_node_ids'] = torch.tensor(var_node_ids) - tensor_dict['variable_node_variable_ids'] = torch.tensor(var_node_variable_ids) - tensor_dict['variable_mention_nums'] = torch.tensor(var_mention_nums, dtype=torch.float) + tensor_dict["variable_node_ids"] = torch.tensor(var_node_ids) + tensor_dict["variable_node_variable_ids"] = torch.tensor( + var_node_variable_ids + ) + tensor_dict["variable_mention_nums"] = torch.tensor( + var_mention_nums, dtype=torch.float + ) return packed_graph, tensor_dict @classmethod - def to_tensor_dict(cls, packed_graph: PackedGraph, - grammar: Grammar, - vocab: Vocab) -> Dict[str, torch.Tensor]: + def to_tensor_dict( + cls, packed_graph: PackedGraph, grammar: Grammar, vocab: Vocab + ) -> Dict[str, torch.Tensor]: obj_name_bpe_model = vocab.obj_name.subtoken_model obj_name_bpe_pad_idx = vocab.obj_name.subtoken_model.pad_id() @@ -337,26 +461,33 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, ast = packed_graph.trees[ast_id] node_type_tokens = [] - if group == 'ast_nodes': + if group == "ast_nodes": node = ast.id_to_node[node_key] type_idx = grammar.syntax_type_to_id[node.node_type] if node.is_variable_node: var_node_name_indices[i] = vocab.source[node.old_name] - # function root with type `block` also has an name entry storing the name of the function - if node.node_type == 'obj' or node.node_type == 'block' and hasattr(node, 'name'): + # function root with type `block` also has an name entry + # storing the name of the function + if ( + node.node_type == "obj" + or node.node_type == "block" + and hasattr(node, "name") + ): # compute variable embedding node_sub_tokens = obj_name_bpe_model.encode_as_ids(node.name) sub_tokens_list.append(node_sub_tokens) node_with_subtokens_indices.append(i) - if hasattr(node, 'type_tokens'): - node_type_tokens = [vocab.grammar.variable_type_to_id(t) for t in node.type_tokens] + if hasattr(node, "type_tokens"): + node_type_tokens = [ + vocab.grammar.variable_type_to_id(t) for t in node.type_tokens + ] - elif group == 'variable_master_nodes': + elif group == "variable_master_nodes": type_idx = variable_master_node_type_idx var_node_name_indices[i] = vocab.source[node_key] - elif group == 'master_nodes': + elif group == "master_nodes": type_idx = master_node_type_idx else: raise ValueError() @@ -368,17 +499,21 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, sub_tokens_indices = None if node_with_subtokens_indices: max_subtoken_num = max(len(x) for x in sub_tokens_list) - sub_tokens_indices = np.zeros((len(sub_tokens_list), max_subtoken_num), dtype=np.int64) + sub_tokens_indices = np.zeros( + (len(sub_tokens_list), max_subtoken_num), dtype=np.int64 + ) sub_tokens_indices.fill(obj_name_bpe_pad_idx) for i, token_ids in enumerate(sub_tokens_list): - sub_tokens_indices[i, :len(token_ids)] = token_ids + sub_tokens_indices[i, : len(token_ids)] = token_ids sub_tokens_indices = torch.from_numpy(sub_tokens_indices) max_typetoken_num = max(len(x) for x in node_type_tokens_list) - node_type_indices = np.zeros((len(node_type_tokens_list), max_typetoken_num), dtype=np.int64) + node_type_indices = np.zeros( + (len(node_type_tokens_list), max_typetoken_num), dtype=np.int64 + ) for i, type_ids in enumerate(node_type_tokens_list): - node_type_indices[i, :len(type_ids)] = type_ids + node_type_indices[i, : len(type_ids)] = type_ids node_type_indices = torch.from_numpy(node_type_indices) return dict( @@ -387,70 +522,116 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, node_syntax_type_indices=node_syntax_type_indices, node_type_indices=node_type_indices, var_node_name_indices=var_node_name_indices, - node_with_subtokens_indices=torch.tensor(node_with_subtokens_indices, dtype=torch.long), + node_with_subtokens_indices=torch.tensor( + node_with_subtokens_indices, dtype=torch.long + ), sub_tokens_indices=sub_tokens_indices, - tree_node_to_tree_id_map=tree_node_to_tree_id_map + tree_node_to_tree_id_map=tree_node_to_tree_id_map, ) - def get_batched_tree_init_encoding(self, tensor_dict: Dict[str, torch.Tensor]) -> torch.Tensor: - node_syntax_type_embedding = self.node_syntax_type_embedding(tensor_dict['node_syntax_type_indices']) - node_type_embedding = self.node_type_embedder(tensor_dict['node_type_indices']) - node_content_embedding = self.var_node_name_embedding(tensor_dict['var_node_name_indices']) * \ - torch.ne(tensor_dict['var_node_name_indices'], 0.).float().unsqueeze(-1) - - if tensor_dict['node_with_subtokens_indices'].size(0) > 0: - obj_node_content_embedding = self.node_content_embedder(tensor_dict['sub_tokens_indices']) - node_content_embedding = node_content_embedding.scatter(0, tensor_dict['node_with_subtokens_indices'].unsqueeze(-1).expand_as(obj_node_content_embedding), - obj_node_content_embedding) - - tree_node_embedding = self.type_and_content_hybrid(torch.cat([node_syntax_type_embedding, node_type_embedding, node_content_embedding], dim=-1)) + def get_batched_tree_init_encoding( + self, tensor_dict: Dict[str, torch.Tensor] + ) -> torch.Tensor: + node_syntax_type_embedding = self.node_syntax_type_embedding( + tensor_dict["node_syntax_type_indices"] + ) + node_type_embedding = self.node_type_embedder(tensor_dict["node_type_indices"]) + node_content_embedding = self.var_node_name_embedding( + tensor_dict["var_node_name_indices"] + ) * torch.ne(tensor_dict["var_node_name_indices"], 0.0).float().unsqueeze(-1) + + if tensor_dict["node_with_subtokens_indices"].size(0) > 0: + obj_node_content_embedding = self.node_content_embedder( + tensor_dict["sub_tokens_indices"] + ) + node_content_embedding = node_content_embedding.scatter( + 0, + tensor_dict["node_with_subtokens_indices"] + .unsqueeze(-1) + .expand_as(obj_node_content_embedding), + obj_node_content_embedding, + ) + + tree_node_embedding = self.type_and_content_hybrid( + torch.cat( + [ + node_syntax_type_embedding, + node_type_embedding, + node_content_embedding, + ], + dim=-1, + ) + ) return tree_node_embedding - def unpack_encoding(self, flattened_node_encodings: torch.Tensor, - batch_syntax_trees: List[AbstractSyntaxTree], - example_node2batch_node_map: Dict): + def unpack_encoding( + self, + flattened_node_encodings: torch.Tensor, + batch_syntax_trees: List[AbstractSyntaxTree], + example_node2batch_node_map: Dict, + ): batch_size = len(batch_syntax_trees) max_node_num = max(tree.size for tree in batch_syntax_trees) index = np.zeros((batch_size, max_node_num), dtype=np.int64) - batch_tree_node_masks = torch.zeros(batch_size, max_node_num, device=self.device) + batch_tree_node_masks = torch.zeros( + batch_size, max_node_num, device=self.device + ) for e_id, syntax_tree in enumerate(batch_syntax_trees): - example_nodes_with_batch_id = [(example_node_id, batch_node_id) - for (_e_id, example_node_id), batch_node_id - in example_node2batch_node_map.items() - if _e_id == e_id] - # example_nodes_batch_id = list(map(lambda x: x[1], sorted(example_nodes_with_batch_id, key=lambda t: t[0]))) - sorted_example_nodes_with_batch_id = sorted(example_nodes_with_batch_id, key=lambda t: t[0]) + example_nodes_with_batch_id = [ + (example_node_id, batch_node_id) + for ( + _e_id, + example_node_id, + ), batch_node_id in example_node2batch_node_map.items() + if _e_id == e_id + ] + sorted_example_nodes_with_batch_id = sorted( + example_nodes_with_batch_id, key=itemgetter(0) + ) example_nodes_batch_id = [t[1] for t in sorted_example_nodes_with_batch_id] - index[e_id, :len(example_nodes_batch_id)] = example_nodes_batch_id - batch_tree_node_masks[e_id, :len(example_nodes_batch_id)] = 1. + index[e_id, : len(example_nodes_batch_id)] = example_nodes_batch_id + batch_tree_node_masks[e_id, : len(example_nodes_batch_id)] = 1.0 # (batch_size, max_node_num, node_encoding_size) - batch_node_encoding = flattened_node_encodings[torch.from_numpy(index).to(flattened_node_encodings.device)] - batch_node_encoding.data.masked_fill_((1. - batch_tree_node_masks).bool().unsqueeze(-1), 0.) + batch_node_encoding = flattened_node_encodings[ + torch.from_numpy(index).to(flattened_node_encodings.device) + ] + batch_node_encoding.data.masked_fill_( + (1.0 - batch_tree_node_masks).bool().unsqueeze(-1), 0.0 + ) - return dict(batch_tree_node_encoding=batch_node_encoding, - batch_tree_node_masks=batch_tree_node_masks) + return dict( + batch_tree_node_encoding=batch_node_encoding, + batch_tree_node_masks=batch_tree_node_masks, + ) def get_decoder_init_state(self, context_encoding, config=None): # compute initial decoder's state via average pooling # (packed_graph_size, encoding_size) - packed_tree_node_encoding = context_encoding['packed_tree_node_encoding'] + packed_tree_node_encoding = context_encoding["packed_tree_node_encoding"] - tree_num = context_encoding['tree_num'] - total_node_num = context_encoding['tree_node_to_tree_id_map'].size(0) + tree_num = context_encoding["tree_num"] + total_node_num = context_encoding["tree_node_to_tree_id_map"].size(0) encoding_size = packed_tree_node_encoding.size(-1) zero_encoding = packed_tree_node_encoding.new_zeros(tree_num, encoding_size) - node_encoding_sum = zero_encoding.scatter_add_(0, context_encoding['tree_node_to_tree_id_map'].unsqueeze( - -1).expand(-1, encoding_size), - packed_tree_node_encoding) - tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_(0, context_encoding['tree_node_to_tree_id_map'], - packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.)) + node_encoding_sum = zero_encoding.scatter_add_( + 0, + context_encoding["tree_node_to_tree_id_map"] + .unsqueeze(-1) + .expand(-1, encoding_size), + packed_tree_node_encoding, + ) + tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_( + 0, + context_encoding["tree_node_to_tree_id_map"], + packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.0), + ) avg_node_encoding = node_encoding_sum / tree_node_num.unsqueeze(-1) c_0 = self.decoder_cell_init(avg_node_encoding) @@ -458,14 +639,17 @@ def get_decoder_init_state(self, context_encoding, config=None): return h_0, c_0 - def get_attention_memory(self, context_encoding, att_target='terminal_nodes'): - if att_target == 'ast_nodes': - memory = context_encoding['packed_tree_node_encoding'][context_encoding['tree_restoration_indices']] - mask = context_encoding['tree_restoration_indices_mask'] - elif att_target == 'terminal_nodes': - memory = context_encoding['packed_tree_node_encoding'][context_encoding['terminal_node_restoration_indices']] - mask = context_encoding['terminal_node_restoration_indices_mask'] + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + packed_tree_node_enc = context_encoding["packed_tree_node_encoding"] + if att_target == "ast_nodes": + memory = packed_tree_node_enc[context_encoding["tree_restoration_indices"]] + mask = context_encoding["tree_restoration_indices_mask"] + elif att_target == "terminal_nodes": + memory = packed_tree_node_enc[ + context_encoding["terminal_node_restoration_indices"] + ] + mask = context_encoding["terminal_node_restoration_indices_mask"] else: - raise ValueError('unknown attention target') + raise ValueError("unknown attention target") return memory, mask diff --git a/neural-model/model/hybrid_encoder.py b/neural-model/model/hybrid_encoder.py index dbad96c..9ba9d5f 100644 --- a/neural-model/model/hybrid_encoder.py +++ b/neural-model/model/hybrid_encoder.py @@ -1,34 +1,29 @@ -from typing import List, Dict, Tuple +from typing import Dict -import numpy as np import torch -from torch import nn as nn - -from model.embedding import NodeTypeEmbedder, SubTokenEmbedder -from model.encoder import Encoder from model.graph_encoder import GraphASTEncoder from model.sequential_encoder import SequentialEncoder -from model.gnn import GatedGraphNeuralNetwork, AdjacencyList +from torch import nn as nn from utils import util -from utils.ast import AbstractSyntaxTree -from utils.grammar import Grammar -from utils.graph import PackedGraph -from utils.vocab import Vocab class HybridEncoder(nn.Module): def __init__(self, config): super(HybridEncoder, self).__init__() - self.graph_encoder = GraphASTEncoder.build(config['graph_encoder']) - self.seq_encoder = SequentialEncoder.build(config['seq_encoder']) - - self.hybrid_method = config['hybrid_method'] - if self.hybrid_method == 'linear_proj': - self.projection = nn.Linear(config['seq_encoder']['source_encoding_size'] + config['graph_encoder']['gnn']['hidden_size'], - config['source_encoding_size'], bias=False) + self.graph_encoder = GraphASTEncoder.build(config["graph_encoder"]) + self.seq_encoder = SequentialEncoder.build(config["seq_encoder"]) + + self.hybrid_method = config["hybrid_method"] + if self.hybrid_method == "linear_proj": + self.projection = nn.Linear( + config["seq_encoder"]["source_encoding_size"] + + config["graph_encoder"]["gnn"]["hidden_size"], + config["source_encoding_size"], + bias=False, + ) else: - assert self.hybrid_method == 'concat' + assert self.hybrid_method == "concat" self.config = config @@ -41,7 +36,7 @@ def default_params(cls): return { "graph_encoder": GraphASTEncoder.default_params(), "seq_encoder": SequentialEncoder.default_params(), - "hybrid_method": "linear_proj" + "hybrid_method": "linear_proj", } @classmethod @@ -51,46 +46,69 @@ def build(cls, config): return cls(params) def forward(self, tensor_dict: Dict[str, torch.Tensor]): - graph_encoding = self.graph_encoder(tensor_dict['graph_encoder_input']) - seq_encoding = self.seq_encoder(tensor_dict['seq_encoder_input']) + graph_encoding = self.graph_encoder(tensor_dict["graph_encoder_input"]) + seq_encoding = self.seq_encoder(tensor_dict["seq_encoder_input"]) - graph_var_encoding = graph_encoding['variable_encoding'] - seq_var_encoding = seq_encoding['variable_encoding'] + graph_var_encoding = graph_encoding["variable_encoding"] + seq_var_encoding = seq_encoding["variable_encoding"] - if self.hybrid_method == 'linear_proj': - variable_encoding = self.projection(torch.cat([graph_var_encoding, seq_var_encoding], dim=-1)) + if self.hybrid_method == "linear_proj": + variable_encoding = self.projection( + torch.cat([graph_var_encoding, seq_var_encoding], dim=-1) + ) else: - variable_encoding = torch.cat([graph_var_encoding, seq_var_encoding], dim=-1) + variable_encoding = torch.cat( + [graph_var_encoding, seq_var_encoding], dim=-1 + ) context_encoding = dict( - batch_size=tensor_dict['batch_size'], + batch_size=tensor_dict["batch_size"], variable_encoding=variable_encoding, graph_encoding_result=graph_encoding, - seq_encoding_result=seq_encoding + seq_encoding_result=seq_encoding, ) return context_encoding def get_decoder_init_state(self, context_encoding, config=None): - gnn_dec_init_state, gnn_dec_init_cell = self.graph_encoder.get_decoder_init_state(context_encoding['graph_encoding_result']) - seq_dec_init_state, seq_dec_init_cell = self.seq_encoder.get_decoder_init_state(context_encoding['seq_encoding_result']) + ( + gnn_dec_init_state, + gnn_dec_init_cell, + ) = self.graph_encoder.get_decoder_init_state( + context_encoding["graph_encoding_result"] + ) + seq_dec_init_state, seq_dec_init_cell = self.seq_encoder.get_decoder_init_state( + context_encoding["seq_encoding_result"] + ) dec_init_state = (gnn_dec_init_state + seq_dec_init_state) / 2 dec_init_cell = (gnn_dec_init_cell + seq_dec_init_cell) / 2 return dec_init_state, dec_init_cell - def get_attention_memory(self, context_encoding, att_target='terminal_nodes'): - assert att_target == 'terminal_nodes' + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + assert att_target == "terminal_nodes" - seq_memory, seq_mask = self.seq_encoder.get_attention_memory(context_encoding['seq_encoding_result'], att_target='terminal_nodes') - gnn_memory, gnn_mask = self.graph_encoder.get_attention_memory(context_encoding['graph_encoding_result'], att_target='terminal_nodes') + seq_memory, seq_mask = self.seq_encoder.get_attention_memory( + context_encoding["seq_encoding_result"], att_target="terminal_nodes" + ) + gnn_memory, gnn_mask = self.graph_encoder.get_attention_memory( + context_encoding["graph_encoding_result"], att_target="terminal_nodes" + ) if gnn_memory.size(-1) < seq_memory.size(-1): # pad values - gnn_memory = torch.cat([gnn_memory.new_zeros(gnn_memory.size(0), gnn_memory.size(1), - seq_memory.size(-1) - gnn_memory.size(-1)), - gnn_memory], dim=-1) + gnn_memory = torch.cat( + [ + gnn_memory.new_zeros( + gnn_memory.size(0), + gnn_memory.size(1), + seq_memory.size(-1) - gnn_memory.size(-1), + ), + gnn_memory, + ], + dim=-1, + ) memory = torch.cat([gnn_memory, seq_memory], dim=1) mask = torch.cat([gnn_mask, seq_mask], dim=1) diff --git a/neural-model/model/model.py b/neural-model/model/model.py index 825a893..76a56ce 100644 --- a/neural-model/model/model.py +++ b/neural-model/model/model.py @@ -1,29 +1,21 @@ import pprint import sys -from typing import List, Dict, Tuple, Iterable +from typing import Dict, Iterable, List, Tuple import torch import torch.nn as nn - -from utils import nn_util, util -from utils.ast import AbstractSyntaxTree from model.decoder import Decoder -from model.recurrent_subtoken_decoder import RecurrentSubtokenDecoder -from model.attentional_recurrent_subtoken_decoder import AttentionalRecurrentSubtokenDecoder -from model.recurrent_decoder import RecurrentDecoder -from model.simple_decoder import SimpleDecoder from model.encoder import Encoder +from model.graph_encoder import GraphASTEncoder from model.hybrid_encoder import HybridEncoder from model.sequential_encoder import SequentialEncoder -from model.graph_encoder import GraphASTEncoder -from utils.graph import PackedGraph +from utils import util from utils.dataset import Batcher, Example -from utils.vocab import SAME_VARIABLE_TOKEN class RenamingModel(nn.Module): def __init__(self, encoder: Encoder, decoder: Decoder): - super(RenamingModel, self).__init__() + super().__init__() self.encoder = encoder self.decoder = decoder @@ -35,9 +27,9 @@ def vocab(self): @property def batcher(self): - if not hasattr(self, '_batcher'): + if not hasattr(self, "_batcher"): _batcher = Batcher(self.config) - setattr(self, '_batcher', _batcher) + setattr(self, "_batcher", _batcher) return self._batcher @@ -45,52 +37,56 @@ def batcher(self): def device(self): return self.encoder.device - @classmethod - def default_params(cls): + @staticmethod + def default_params(): return { - 'train': { - 'unchanged_variable_weight': 1.0, - 'max_epoch': 30, - 'patience': 5 - }, - 'decoder': { - 'type': 'SimpleDecoder' - } + "train": {"unchanged_variable_weight": 1.0, "max_epoch": 30, "patience": 5}, + "decoder": {"type": "SimpleDecoder"}, } @classmethod - def build(cls, config): + def build(cls, config, logging=False): params = util.update(cls.default_params(), config) - encoder = globals()[config['encoder']['type']].build(config['encoder']) - decoder = globals()[config['decoder']['type']].build(config['decoder']) + + encoder = globals()[config["encoder"]["type"]].build(config["encoder"]) + decoder = globals()[config["decoder"]["type"]].build(config["decoder"]) model = cls(encoder, decoder) - params = util.update(params, {'encoder': encoder.config, - 'decoder': decoder.config}) + params = util.update( + params, {"encoder": encoder.config, "decoder": decoder.config} + ) model.config = params - model.decoder.encoder = encoder # give the decoder a reference to the encoder + # give the decoder a reference to the encoder + model.decoder.encoder = encoder # assign batcher to sub-modules encoder.batcher = model.batcher decoder.batcher = model.batcher - print('Current Configuration:', file=sys.stderr) - pp = pprint.PrettyPrinter(indent=2, stream=sys.stderr) - pp.pprint(model.config) + if logging: + print("Current Configuration:", file=sys.stderr) + pp = pprint.PrettyPrinter(indent=2, stream=sys.stderr) + pp.pprint(model.config) return model - def forward(self, source_asts: Dict[str, torch.Tensor], prediction_target: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, Dict]: + def forward( + self, + source_asts: Dict[str, torch.Tensor], + prediction_target: Dict[str, torch.Tensor], + ) -> Tuple[torch.Tensor, Dict]: """ - Given a batch of decompiled abstract syntax trees, and the gold-standard renaming of variable nodes, - compute the log-likelihood of the gold-standard renaming for training + Given a batch of decompiled abstract syntax trees, and the + gold-standard renaming of variable nodes, compute the log-likelihood of + the gold-standard renaming for training Arg: source_asts: a list of ASTs - variable_name_maps: mapping of decompiled variable names to its renamed values + variable_name_maps: mapping of decompiled variable names to its + renamed values Return: - a tensor of size (batch_size) denoting the log-likelihood of renamings + tensor of size batch_size denoting the log-likelihood of renamings """ # src_ast_encoding: (batch_size, max_ast_node_num, node_encoding_size) @@ -100,54 +96,62 @@ def forward(self, source_asts: Dict[str, torch.Tensor], prediction_target: Dict[ # (batch_size, variable_num, vocab_size) or (prediction_node_num, vocab_size) var_name_log_probs = self.decoder(context_encoding, prediction_target) - result = self.decoder.get_target_log_prob(var_name_log_probs, prediction_target, context_encoding) - - tgt_var_name_log_prob = result['tgt_var_name_log_prob'] - tgt_weight = prediction_target['variable_tgt_name_weight'] + result = self.decoder.get_target_log_prob( + var_name_log_probs, prediction_target, context_encoding + ) + tgt_var_name_log_prob = result["tgt_var_name_log_prob"] + tgt_weight = prediction_target["variable_tgt_name_weight"] weighted_log_prob = tgt_var_name_log_prob * tgt_weight - - ast_log_probs = weighted_log_prob.sum(dim=-1) / prediction_target['target_variable_encoding_indices_mask'].sum(-1) - result['batch_log_prob'] = ast_log_probs + tgt_var_encoding_indices_mask = prediction_target[ + "target_variable_encoding_indices_mask" + ] + ast_log_probs = weighted_log_prob.sum( + dim=-1 + ) / tgt_var_encoding_indices_mask.sum(-1) + result["batch_log_prob"] = ast_log_probs return result - def decode_dataset(self, dataset, batch_size=4096) -> Iterable[Tuple[Example, Dict]]: + def decode_dataset( + self, dataset, batch_size=4096 + ) -> Iterable[Tuple[Example, Dict]]: with torch.no_grad(): - data_iter = dataset.batch_iterator(batch_size=batch_size, train=False, progress=False, - config=self.config) - was_training = self.training + data_iter = dataset.batch_iterator( + batch_size=batch_size, train=False, progress=False, config=self.config + ) self.eval() for batch in data_iter: - rename_results = self.decoder.predict([e.ast for e in batch.examples], self.encoder) - for example, rename_result in zip(batch.examples, rename_results): - yield example, rename_result + rename_results = self.decoder.predict( + [e.ast for e in batch.examples], self.encoder + ) + + for example, result in zip(batch.examples, rename_results): + yield example, result def predict(self, examples: List[Example]): return self.decoder.predict(examples, self.encoder) def save(self, model_path, **kwargs): params = { - 'config': self.config, - 'state_dict': self.state_dict(), - 'kwargs': kwargs + "config": self.config, + "state_dict": self.state_dict(), + "kwargs": kwargs, } torch.save(params, model_path) @classmethod - def load(cls, model_path, use_cuda=False, new_config=None) -> 'RenamingModel': + def load(cls, model_path, use_cuda=False, new_config=None) -> "RenamingModel": device = torch.device("cuda:0" if use_cuda else "cpu") - params = torch.load(model_path, map_location=lambda storage, loc: storage) - config = params['config'] + params = torch.load(model_path, map_location=lambda storage, _loc: storage) - if new_config: - config = util.update(config, new_config) + config = util.update(params["config"], new_config) - kwargs = params['kwargs'] if params['kwargs'] is not None else dict() + kwargs = dict() if params["kwargs"] is None else params["kwargs"] model = cls.build(config, **kwargs) - model.load_state_dict(params['state_dict'], strict=False) + model.load_state_dict(params["state_dict"], strict=False) model = model.to(device) model.eval() diff --git a/neural-model/model/recurrent_decoder.py b/neural-model/model/recurrent_decoder.py index 5c2eb46..a51ce84 100644 --- a/neural-model/model/recurrent_decoder.py +++ b/neural-model/model/recurrent_decoder.py @@ -1,29 +1,38 @@ -from collections import namedtuple, OrderedDict +from collections import OrderedDict, namedtuple from typing import Dict, List import torch -from torch import nn as nn - from model.decoder import Decoder from model.encoder import Encoder -from utils import util, nn_util +from torch import nn as nn +from utils import nn_util, util from utils.ast import AbstractSyntaxTree -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab class RecurrentDecoder(Decoder): - def __init__(self, ast_node_encoding_size: int, hidden_size: int, dropout: float, vocab: Vocab): + def __init__( + self, + ast_node_encoding_size: int, + hidden_size: int, + dropout: float, + vocab: Vocab, + ): super(Decoder, self).__init__() self.vocab = vocab - self.lstm_cell = nn.LSTMCell(ast_node_encoding_size + ast_node_encoding_size, hidden_size) + self.lstm_cell = nn.LSTMCell( + ast_node_encoding_size + ast_node_encoding_size, hidden_size + ) self.decoder_cell_init = nn.Linear(ast_node_encoding_size, hidden_size) - self.state2names = nn.Linear(ast_node_encoding_size, len(vocab.target), bias=True) + self.state2names = nn.Linear( + ast_node_encoding_size, len(vocab.target), bias=True + ) self.dropout = nn.Dropout(dropout) self.config: Dict = None - self.Hypothesis = namedtuple('Hypothesis', ['variable_list', 'score']) + self.Hypothesis = namedtuple("Hypothesis", ["variable_list", "score"]) @property def device(self): @@ -32,21 +41,26 @@ def device(self): @classmethod def default_params(cls): return { - 'vocab_file': None, - 'ast_node_encoding_size': 128, - 'hidden_size': 128, - 'input_feed': False, - 'dropout': 0.2, - 'beam_size': 5, - 'unk_replace': True + "vocab_file": None, + "ast_node_encoding_size": 128, + "hidden_size": 128, + "input_feed": False, + "dropout": 0.2, + "beam_size": 5, + "unk_replace": True, } @classmethod def build(cls, config): params = util.update(cls.default_params(), config) - vocab = Vocab.load(params['vocab_file']) - model = cls(params['ast_node_encoding_size'], params['hidden_size'], params['dropout'], vocab) + vocab = Vocab.load(params["vocab_file"]) + model = cls( + params["ast_node_encoding_size"], + params["hidden_size"], + params["dropout"], + vocab, + ) model.config = params return model @@ -55,19 +69,26 @@ def get_init_state(self, src_ast_encoding): # compute initial decoder's state via average pooling # (packed_graph_size, encoding_size) - packed_tree_node_encoding = src_ast_encoding['packed_tree_node_encoding'] + packed_tree_node_encoding = src_ast_encoding["packed_tree_node_encoding"] - tree_num = src_ast_encoding['tree_num'] - total_node_num = src_ast_encoding['tree_node_to_tree_id_map'].size(0) + tree_num = src_ast_encoding["tree_num"] + total_node_num = src_ast_encoding["tree_node_to_tree_id_map"].size(0) encoding_size = packed_tree_node_encoding.size(-1) zero_encoding = packed_tree_node_encoding.new_zeros(tree_num, encoding_size) - node_encoding_sum = zero_encoding.scatter_add_(0, - src_ast_encoding['tree_node_to_tree_id_map'].unsqueeze(-1).expand(-1, encoding_size), - packed_tree_node_encoding) - tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_(0, src_ast_encoding['tree_node_to_tree_id_map'], - packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.)) + node_encoding_sum = zero_encoding.scatter_add_( + 0, + src_ast_encoding["tree_node_to_tree_id_map"] + .unsqueeze(-1) + .expand(-1, encoding_size), + packed_tree_node_encoding, + ) + tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_( + 0, + src_ast_encoding["tree_node_to_tree_id_map"], + packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.0), + ) avg_node_encoding = node_encoding_sum / tree_node_num.unsqueeze(-1) c_0 = self.decoder_cell_init(avg_node_encoding) @@ -85,24 +106,38 @@ def rnn_step(self, x, h_tm1, src_ast_encoding): def forward(self, src_ast_encoding, prediction_target): # (prediction_node_num, encoding_size) - variable_master_node_encoding = src_ast_encoding['variable_master_node_encoding'] + variable_master_node_encoding = src_ast_encoding[ + "variable_master_node_encoding" + ] # (batch_size, max_prediction_node_num) - variable_master_node_restoration_indices = src_ast_encoding['variable_encoding_restoration_indices'] - variable_master_node_restoration_indices_mask = src_ast_encoding['variable_encoding_restoration_indices_mask'] + variable_master_node_restoration_indices = src_ast_encoding[ + "variable_encoding_restoration_indices" + ] + variable_master_node_restoration_indices_mask = src_ast_encoding[ + "variable_encoding_restoration_indices_mask" + ] # (batch_size, max_prediction_node_num, encoding_size) - variable_master_node_encoding = variable_master_node_encoding[variable_master_node_restoration_indices] - variable_tgt_name_id = prediction_target['variable_tgt_name_id'][variable_master_node_restoration_indices] + variable_master_node_encoding = variable_master_node_encoding[ + variable_master_node_restoration_indices + ] + variable_tgt_name_id = prediction_target["variable_tgt_name_id"][ + variable_master_node_restoration_indices + ] batch_size = variable_tgt_name_id.size(0) variable_encoding_size = variable_master_node_encoding.size(-1) h_0 = self.get_init_state(src_ast_encoding) - att_tm1 = variable_master_node_encoding.new_zeros(src_ast_encoding['tree_num'], variable_master_node_encoding.size(-1)) + att_tm1 = variable_master_node_encoding.new_zeros( + src_ast_encoding["tree_num"], variable_master_node_encoding.size(-1) + ) h_tm1 = h_0 query_vecs = [] - for t, variable_encoding_t in enumerate(variable_master_node_encoding.split(split_size=1, dim=1)): + for t, variable_encoding_t in enumerate( + variable_master_node_encoding.split(split_size=1, dim=1) + ): # variable_encoding_t: (batch_size, encoding_size) variable_encoding_t = variable_encoding_t.squeeze(1) if t > 0: @@ -112,9 +147,11 @@ def forward(self, src_ast_encoding, prediction_target): v_tm1_name_embed = self.state2names.weight[v_tm1_name_id] else: # (batch_size, encoding_size) - v_tm1_name_embed = torch.zeros(batch_size, variable_encoding_size, device=self.device) + v_tm1_name_embed = torch.zeros( + batch_size, variable_encoding_size, device=self.device + ) - if self.config['input_feed']: + if self.config["input_feed"]: x = torch.cat([variable_encoding_t, v_tm1_name_embed, att_tm1], dim=-1) else: x = torch.cat([variable_encoding_t, v_tm1_name_embed], dim=-1) @@ -131,31 +168,47 @@ def forward(self, src_ast_encoding, prediction_target): # (batch_size, max_prediction_node_num, vocab_size) logits = self.state2names(query_vecs) var_name_log_probs = torch.log_softmax(logits, dim=-1) - var_name_log_probs = var_name_log_probs * variable_master_node_restoration_indices_mask.unsqueeze(-1) + var_name_log_probs = ( + var_name_log_probs + * variable_master_node_restoration_indices_mask.unsqueeze(-1) + ) return var_name_log_probs - def get_target_log_prob(self, var_name_log_probs, prediction_target, src_ast_encoding): + def get_target_log_prob( + self, var_name_log_probs, prediction_target, src_ast_encoding + ): # (batch_size, max_prediction_node_num) - variable_tgt_name_id = prediction_target['variable_tgt_name_id'][src_ast_encoding['variable_encoding_restoration_indices']] - tgt_var_name_log_prob = torch.gather(var_name_log_probs, - dim=-1, index=variable_tgt_name_id.unsqueeze(-1)).squeeze(-1) - - tgt_var_name_log_prob = tgt_var_name_log_prob * src_ast_encoding['variable_encoding_restoration_indices_mask'] + variable_tgt_name_id = prediction_target["variable_tgt_name_id"][ + src_ast_encoding["variable_encoding_restoration_indices"] + ] + tgt_var_name_log_prob = torch.gather( + var_name_log_probs, dim=-1, index=variable_tgt_name_id.unsqueeze(-1) + ).squeeze(-1) + + tgt_var_name_log_prob = ( + tgt_var_name_log_prob + * src_ast_encoding["variable_encoding_restoration_indices_mask"] + ) result = dict(tgt_var_name_log_prob=tgt_var_name_log_prob) return result - def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> List[Dict]: - beam_size = self.config['beam_size'] - unk_replace = self.config['unk_replace'] + def predict( + self, source_asts: List[AbstractSyntaxTree], encoder: Encoder + ) -> List[Dict]: + beam_size = self.config["beam_size"] + unk_replace = self.config["unk_replace"] variable_nums = [] for ast_id, ast in enumerate(source_asts): variable_nums.append(len(ast.variables)) - beams = OrderedDict((ast_id, [self.Hypothesis([], 0.)]) for ast_id, ast in enumerate(source_asts)) + beams = OrderedDict( + (ast_id, [self.Hypothesis([], 0.0)]) + for ast_id, ast in enumerate(source_asts) + ) hyp_scores_tm1 = torch.zeros(len(beams), 1, device=self.device) completed_hyps = [[] for _ in source_asts] tgt_vocab_size = len(self.vocab.target) @@ -167,16 +220,24 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li h_tm1 = h_0 = self.get_init_state(context_encoding) # (prediction_node_num, encoding_size) - variable_master_node_encoding = context_encoding['variable_master_node_encoding'] + variable_master_node_encoding = context_encoding[ + "variable_master_node_encoding" + ] encoding_size = variable_master_node_encoding.size(1) # (batch_size, max_prediction_node_num) - variable_master_node_restoration_indices = context_encoding['variable_encoding_restoration_indices'] + variable_master_node_restoration_indices = context_encoding[ + "variable_encoding_restoration_indices" + ] # (batch_size, max_prediction_node_num, encoding_size) - variable_master_node_encoding = variable_master_node_encoding[variable_master_node_restoration_indices] + variable_master_node_encoding = variable_master_node_encoding[ + variable_master_node_restoration_indices + ] variable_encoding_t = variable_master_node_encoding[:, 0] # (batch_size, encoding_size) - variable_name_embed_tm1 = att_tm1 = torch.zeros(len(source_asts), encoding_size, device=self.device) + variable_name_embed_tm1 = att_tm1 = torch.zeros( + len(source_asts), encoding_size, device=self.device + ) max_prediction_node_num = variable_master_node_encoding.size(1) @@ -185,11 +246,18 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li live_tree_ids = [ast_id for ast_id in beams] if t > 0: - variable_encoding_t = variable_master_node_encoding[live_tree_ids][:, t]\ - .unsqueeze(1).expand(-1, beam_size, -1).contiguous().view(-1, encoding_size) - - if self.config['input_feed']: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1) + variable_encoding_t = ( + variable_master_node_encoding[live_tree_ids][:, t] + .unsqueeze(1) + .expand(-1, beam_size, -1) + .contiguous() + .view(-1, encoding_size) + ) + + if self.config["input_feed"]: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 + ) else: x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) @@ -198,19 +266,23 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li # (live_beam_num, beam_size, encoding_size) q_t_by_beam = q_t.view(len(beams), -1, q_t.size(-1)) # (live_beam_num, beam_size, vocab_size) - hyp_var_name_scores_t = torch.log_softmax(self.state2names(q_t_by_beam), dim=-1) + hyp_var_name_scores_t = torch.log_softmax( + self.state2names(q_t_by_beam), dim=-1 + ) if unk_replace: - hyp_var_name_scores_t[:, :, self.vocab.target['']] = float('-inf') + hyp_var_name_scores_t[:, :, self.vocab.target[""]] = float("-inf") cont_cand_hyp_scores = hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t # (live_beam_num, beam_size) - new_hyp_scores, new_hyp_position_list = torch.topk(cont_cand_hyp_scores.view(len(beams), -1), k=beam_size, dim=-1) + new_hyp_scores, new_hyp_position_list = torch.topk( + cont_cand_hyp_scores.view(len(beams), -1), k=beam_size, dim=-1 + ) # (live_beam_num, beam_size) - prev_hyp_ids = (new_hyp_position_list / tgt_vocab_size) - hyp_var_name_ids = (new_hyp_position_list % tgt_vocab_size) + prev_hyp_ids = new_hyp_position_list / tgt_vocab_size + hyp_var_name_ids = new_hyp_position_list % tgt_vocab_size new_hyp_scores = new_hyp_scores # move this tensor to cpu for fast indexing @@ -228,8 +300,10 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li hyp_var_name_id = _hyp_var_name_ids[beam_id, i].item() new_hyp_score = _new_hyp_scores[beam_id, i].item() - new_hyp = self.Hypothesis(variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], + score=new_hyp_score, + ) new_hyps.append(new_hyp) if t + 1 == variable_nums[ast_id]: completed_hyps[ast_id] = new_hyps @@ -239,7 +313,9 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li if t < max_prediction_node_num - 1: # (live_beam_num, beam_size, *) - prev_hyp_ids = (torch.arange(len(beams)).to(self.device) * live_beam_size).unsqueeze(-1) + prev_hyp_ids + prev_hyp_ids = ( + torch.arange(len(beams)).to(self.device) * live_beam_size + ).unsqueeze(-1) + prev_hyp_ids live_prev_hyp_ids = prev_hyp_ids[live_beam_ids].view(-1) live_beam_ids = torch.tensor(live_beam_ids, device=self.device) @@ -265,9 +341,11 @@ def predict(self, source_asts: List[AbstractSyntaxTree], encoder: Encoder) -> Li if new_var_name == SAME_VARIABLE_TOKEN: new_var_name = old_name - variable_rename_result[old_name] = {'new_name': new_var_name, - 'prob': top_hyp.score} + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } variable_rename_results.append(variable_rename_result) - return variable_rename_results \ No newline at end of file + return variable_rename_results diff --git a/neural-model/model/recurrent_subtoken_decoder.py b/neural-model/model/recurrent_subtoken_decoder.py index d000d5a..ebf979b 100644 --- a/neural-model/model/recurrent_subtoken_decoder.py +++ b/neural-model/model/recurrent_subtoken_decoder.py @@ -1,20 +1,27 @@ import sys -from collections import namedtuple, OrderedDict +from collections import OrderedDict, namedtuple from typing import Dict, List import torch -from torch import nn as nn - from model.decoder import Decoder from model.encoder import Encoder -from utils import util, nn_util +from torch import nn as nn +from utils import nn_util, util from utils.ast import AbstractSyntaxTree from utils.dataset import Example -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN, END_OF_VARIABLE_TOKEN +from utils.vocab import END_OF_VARIABLE_TOKEN, SAME_VARIABLE_TOKEN, Vocab class RecurrentSubtokenDecoder(Decoder): - def __init__(self, variable_encoding_size: int, hidden_size: int, dropout: float, tie_embed: bool, input_feed: bool, vocab: Vocab): + def __init__( + self, + variable_encoding_size: int, + hidden_size: int, + dropout: float, + tie_embed: bool, + input_feed: bool, + vocab: Vocab, + ): super(Decoder, self).__init__() self.vocab = vocab @@ -31,7 +38,9 @@ def __init__(self, variable_encoding_size: int, hidden_size: int, dropout: float self.dropout = nn.Dropout(dropout) self.config: Dict = None - self.Hypothesis = namedtuple('Hypothesis', ['variable_list', 'variable_ptr', 'score']) + self.Hypothesis = namedtuple( + "Hypothesis", ["variable_list", "variable_ptr", "score"] + ) @property def device(self): @@ -40,28 +49,34 @@ def device(self): @classmethod def default_params(cls): return { - 'vocab_file': None, - 'variable_encoding_size': 128, - 'hidden_size': 128, - 'input_feed': False, - 'tie_embedding': True, - 'dropout': 0.2, - 'beam_size': 5, - 'max_prediction_time_step': 1200, - 'independent_prediction_for_each_variable': False + "vocab_file": None, + "variable_encoding_size": 128, + "hidden_size": 128, + "input_feed": False, + "tie_embedding": True, + "dropout": 0.2, + "beam_size": 5, + "max_prediction_time_step": 1200, + "independent_prediction_for_each_variable": False, } @property def independent_prediction_for_each_variable(self): - return self.config['independent_prediction_for_each_variable'] + return self.config["independent_prediction_for_each_variable"] @classmethod def build(cls, config): params = util.update(cls.default_params(), config) - vocab = Vocab.load(params['vocab_file']) - model = cls(params['variable_encoding_size'], - params['hidden_size'], params['dropout'], params['tie_embedding'], params['input_feed'], vocab) + vocab = Vocab.load(params["vocab_file"]) + model = cls( + params["variable_encoding_size"], + params["hidden_size"], + params["dropout"], + params["tie_embedding"], + params["input_feed"], + vocab, + ) model.config = params return model @@ -79,31 +94,46 @@ def rnn_step(self, x, h_tm1, src_ast_encoding): def forward(self, src_ast_encoding, prediction_target): # (batch_size, max_time_step) - target_variable_encoding_indices = prediction_target['target_variable_encoding_indices'] - target_variable_encoding_indices_mask = prediction_target['target_variable_encoding_indices_mask'] + target_variable_encoding_indices = prediction_target[ + "target_variable_encoding_indices" + ] + target_variable_encoding_indices_mask = prediction_target[ + "target_variable_encoding_indices_mask" + ] batch_size = target_variable_encoding_indices.size(0) - variable_encoding_size = src_ast_encoding['variable_encoding'].size(-1) + variable_encoding_size = src_ast_encoding["variable_encoding"].size(-1) # (batch_size, max_time_step, encoding_size) # scatter variable encoding to sub-token time steps - variable_encoding = torch.gather(src_ast_encoding['variable_encoding'], 1, - target_variable_encoding_indices.unsqueeze(-1).expand(-1, -1, variable_encoding_size)) + variable_encoding = torch.gather( + src_ast_encoding["variable_encoding"], + 1, + target_variable_encoding_indices.unsqueeze(-1).expand( + -1, -1, variable_encoding_size + ), + ) # (batch_size, max_time_step, encoding_size) - variable_tgt_name_id = prediction_target['variable_tgt_name_id'] + variable_tgt_name_id = prediction_target["variable_tgt_name_id"] h_0 = self.get_init_state(src_ast_encoding) - att_tm1 = variable_encoding.new_zeros(src_ast_encoding['batch_size'], self.lstm_cell.hidden_size) - v_tm1_name_embed = torch.zeros(batch_size, self.lstm_cell.hidden_size, device=self.device) + att_tm1 = variable_encoding.new_zeros( + src_ast_encoding["batch_size"], self.lstm_cell.hidden_size + ) + v_tm1_name_embed = torch.zeros( + batch_size, self.lstm_cell.hidden_size, device=self.device + ) h_tm1 = h_0 query_vecs = [] max_time_step = variable_encoding.size(1) - for t, variable_encoding_t in enumerate(variable_encoding.split(split_size=1, dim=1)): + for t, variable_encoding_t in enumerate( + variable_encoding.split(split_size=1, dim=1) + ): # variable_encoding_t: (batch_size, encoding_size) variable_encoding_t = variable_encoding_t.squeeze(1) - if self.config['input_feed']: + if self.config["input_feed"]: x = torch.cat([variable_encoding_t, v_tm1_name_embed, att_tm1], dim=-1) else: x = torch.cat([variable_encoding_t, v_tm1_name_embed], dim=-1) @@ -114,7 +144,7 @@ def forward(self, src_ast_encoding, prediction_target): h_tm1 = h_t query_vecs.append(q_t) v_tm1_name_id = variable_tgt_name_id[:, t] - if self.config['tie_embedding']: + if self.config["tie_embedding"]: v_tm1_name_embed = self.state2names.weight[v_tm1_name_id] else: v_tm1_name_embed = self.var_name_embed(v_tm1_name_id) @@ -124,8 +154,13 @@ def forward(self, src_ast_encoding, prediction_target): variable_ids_tp1 = target_variable_encoding_indices[:, t + 1] variable_ids_t = target_variable_encoding_indices[:, t] - is_tp1_same_variable = torch.eq(variable_ids_tp1, variable_ids_t).float().unsqueeze(-1) # TODO: check if correct! - h_tm1 = (h_tm1[0] * is_tp1_same_variable, h_tm1[1] * is_tp1_same_variable) + is_tp1_same_variable = ( + torch.eq(variable_ids_tp1, variable_ids_t).float().unsqueeze(-1) + ) # TODO: check if correct! + h_tm1 = ( + h_tm1[0] * is_tp1_same_variable, + h_tm1[1] * is_tp1_same_variable, + ) att_tm1 = att_tm1 * is_tp1_same_variable v_tm1_name_embed = v_tm1_name_embed * is_tp1_same_variable @@ -135,17 +170,25 @@ def forward(self, src_ast_encoding, prediction_target): # (batch_size, max_prediction_node_num, vocab_size) logits = self.state2names(query_vecs) var_name_log_probs = torch.log_softmax(logits, dim=-1) - var_name_log_probs = var_name_log_probs * target_variable_encoding_indices_mask.unsqueeze(-1) + var_name_log_probs = ( + var_name_log_probs * target_variable_encoding_indices_mask.unsqueeze(-1) + ) return var_name_log_probs - def get_target_log_prob(self, var_name_log_probs, prediction_target, src_ast_encoding): + def get_target_log_prob( + self, var_name_log_probs, prediction_target, src_ast_encoding + ): # (batch_size, max_prediction_node_num) - variable_tgt_name_id = prediction_target['variable_tgt_name_id'] - tgt_var_name_log_prob = torch.gather(var_name_log_probs, - dim=-1, index=variable_tgt_name_id.unsqueeze(-1)).squeeze(-1) + variable_tgt_name_id = prediction_target["variable_tgt_name_id"] + tgt_var_name_log_prob = torch.gather( + var_name_log_probs, dim=-1, index=variable_tgt_name_id.unsqueeze(-1) + ).squeeze(-1) - tgt_var_name_log_prob = tgt_var_name_log_prob * prediction_target['target_variable_encoding_indices_mask'] + tgt_var_name_log_prob = ( + tgt_var_name_log_prob + * prediction_target["target_variable_encoding_indices_mask"] + ) result = dict(tgt_var_name_log_prob=tgt_var_name_log_prob) @@ -153,7 +196,7 @@ def get_target_log_prob(self, var_name_log_probs, prediction_target, src_ast_enc def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: batch_size = len(examples) - beam_size = self.config['beam_size'] + beam_size = self.config["beam_size"] same_variable_id = self.vocab.target[SAME_VARIABLE_TOKEN] end_of_variable_id = self.vocab.target[END_OF_VARIABLE_TOKEN] @@ -161,7 +204,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: for ast_id, example in enumerate(examples): variable_nums.append(len(example.ast.variables)) - beams = OrderedDict((ast_id, [self.Hypothesis([], 0, 0.)]) for ast_id in range(batch_size)) + beams = OrderedDict( + (ast_id, [self.Hypothesis([], 0, 0.0)]) for ast_id in range(batch_size) + ) hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) completed_hyps = [[] for _ in range(batch_size)] tgt_vocab_size = len(self.vocab.target) @@ -174,20 +219,26 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: # Note that we are using the `restoration_indices` from `context_encoding`, which is the word-level restoration index # (batch_size, variable_master_node_num, encoding_size) - variable_encoding = context_encoding['variable_encoding'] + variable_encoding = context_encoding["variable_encoding"] # (batch_size, encoding_size) - variable_name_embed_tm1 = att_tm1 = torch.zeros(batch_size, self.lstm_cell.hidden_size, device=self.device) + variable_name_embed_tm1 = att_tm1 = torch.zeros( + batch_size, self.lstm_cell.hidden_size, device=self.device + ) - max_prediction_time_step = self.config['max_prediction_time_step'] + max_prediction_time_step = self.config["max_prediction_time_step"] for t in range(0, max_prediction_time_step): # (total_live_hyp_num, encoding_size) if t > 0: - variable_encoding_t = variable_encoding[hyp_ast_ids_t, hyp_variable_ptrs_t] + variable_encoding_t = variable_encoding[ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] else: variable_encoding_t = variable_encoding[:, 0] - if self.config['input_feed']: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1) + if self.config["input_feed"]: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 + ) else: x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) @@ -210,11 +261,13 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: for beam_id, (ast_id, beam) in enumerate(beams.items()): beam_end_hyp_pos = beam_start_hyp_pos + len(beam) # (live_beam_size, vocab_size) - beam_cont_cand_hyp_scores = cont_cand_hyp_scores[beam_start_hyp_pos: beam_end_hyp_pos] + beam_cont_cand_hyp_scores = cont_cand_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] cont_beam_size = beam_size - len(completed_hyps[ast_id]) - beam_new_hyp_scores, beam_new_hyp_positions = torch.topk(beam_cont_cand_hyp_scores.view(-1), - k=cont_beam_size, - dim=-1) + beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( + beam_cont_cand_hyp_scores.view(-1), k=cont_beam_size, dim=-1 + ) # (cont_beam_size) beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size @@ -235,12 +288,17 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: variable_ptr += 1 # remove empty cases - if len(prev_hyp.variable_list) == 0 or prev_hyp.variable_list[-1] == end_of_variable_id: + if ( + len(prev_hyp.variable_list) == 0 + or prev_hyp.variable_list[-1] == end_of_variable_id + ): continue - new_hyp = self.Hypothesis(variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], - variable_ptr=variable_ptr, - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], + variable_ptr=variable_ptr, + score=new_hyp_score, + ) if variable_ptr == variable_nums[ast_id]: completed_hyps[ast_id].append(new_hyp) @@ -252,7 +310,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: new_hyp_var_name_ids.append(hyp_var_name_id) new_hyp_ast_ids.append(ast_id) new_hyp_variable_ptrs.append(variable_ptr) - is_same_variable_mask.append(1. if prev_hyp.variable_ptr == variable_ptr else 0.) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) beam_start_hyp_pos = beam_end_hyp_pos @@ -268,10 +328,17 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: beams = new_beams if self.independent_prediction_for_each_variable: - is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) - h_tm1 = (h_tm1[0] * is_same_variable_mask, h_tm1[1] * is_same_variable_mask) + is_same_variable_mask = torch.tensor( + is_same_variable_mask, device=self.device, dtype=torch.float + ).unsqueeze(-1) + h_tm1 = ( + h_tm1[0] * is_same_variable_mask, + h_tm1[1] * is_same_variable_mask, + ) att_tm1 = att_tm1 * is_same_variable_mask - variable_name_embed_tm1 = variable_name_embed_tm1 * is_same_variable_mask + variable_name_embed_tm1 = ( + variable_name_embed_tm1 * is_same_variable_mask + ) else: break @@ -283,10 +350,15 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: if not hyps: # return identity renamings - print(f'Failed to found a hypothesis for function {ast.compilation_unit}', file=sys.stderr) + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) for old_name in ast.variables: - variable_rename_result[old_name] = {'new_name': old_name, - 'prob': 0.} + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } else: top_hyp = hyps[0] sub_token_ptr = 0 @@ -297,14 +369,20 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: sub_token_ptr += 1 # point to first sub-token of next variable sub_token_end = sub_token_ptr - var_name_token_ids = top_hyp.variable_list[sub_token_begin: sub_token_end] # include ending + var_name_token_ids = top_hyp.variable_list[ + sub_token_begin:sub_token_end + ] # include ending if var_name_token_ids == [same_variable_id, end_of_variable_id]: new_var_name = old_name else: - new_var_name = self.vocab.target.subtoken_model.decode_ids(var_name_token_ids) - - variable_rename_result[old_name] = {'new_name': new_var_name, - 'prob': top_hyp.score} + new_var_name = self.vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } variable_rename_results.append(variable_rename_result) diff --git a/neural-model/model/sequential_encoder.py b/neural-model/model/sequential_encoder.py index db00f85..a8f0db6 100644 --- a/neural-model/model/sequential_encoder.py +++ b/neural-model/model/sequential_encoder.py @@ -1,31 +1,39 @@ import sys from typing import Dict, List -import numpy as np - -from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence +import numpy as np +import torch +import torch.nn as nn from model.encoder import * +from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence from utils import nn_util, util from utils.dataset import Example from utils.vocab import PAD_ID, Vocab -import torch -import torch.nn as nn class SequentialEncoder(Encoder): def __init__(self, config): super().__init__() - self.vocab = vocab = Vocab.load(config['vocab_file']) + self.vocab = vocab = Vocab.load(config["vocab_file"]) - self.src_word_embed = nn.Embedding(len(vocab.source_tokens), config['source_embedding_size']) + self.src_word_embed = nn.Embedding( + len(vocab.source_tokens), config["source_embedding_size"] + ) - dropout = config['dropout'] - self.lstm_encoder = nn.LSTM(input_size=self.src_word_embed.embedding_dim, - hidden_size=config['source_encoding_size'] // 2, num_layers=config['num_layers'], - batch_first=True, bidirectional=True, dropout=dropout) + dropout = config["dropout"] + self.lstm_encoder = nn.LSTM( + input_size=self.src_word_embed.embedding_dim, + hidden_size=config["source_encoding_size"] // 2, + num_layers=config["num_layers"], + batch_first=True, + bidirectional=True, + dropout=dropout, + ) - self.decoder_cell_init = nn.Linear(config['source_encoding_size'], config['decoder_hidden_size']) + self.decoder_cell_init = nn.Linear( + config["source_encoding_size"], config["decoder_hidden_size"] + ) self.dropout = nn.Dropout(dropout) self.config = config @@ -37,11 +45,11 @@ def device(self): @classmethod def default_params(cls): return { - 'source_encoding_size': 256, - 'decoder_hidden_size': 128, - 'source_embedding_size': 128, - 'vocab_file': None, - 'num_layers': 1 + "source_encoding_size": 256, + "decoder_hidden_size": 128, + "source_embedding_size": 128, + "vocab_file": None, + "num_layers": 1, } @classmethod @@ -51,16 +59,20 @@ def build(cls, config): return cls(params) def forward(self, tensor_dict: Dict[str, torch.Tensor]): - code_token_encoding, code_token_mask, (last_states, last_cells) = self.encode_sequence(tensor_dict['src_code_tokens']) + ( + code_token_encoding, + code_token_mask, + (last_states, last_cells), + ) = self.encode_sequence(tensor_dict["src_code_tokens"]) # (batch_size, max_variable_mention_num) # variable_mention_positions = tensor_dict['variable_position'] - variable_mention_mask = tensor_dict['variable_mention_mask'] - variable_mention_to_variable_id = tensor_dict['variable_mention_to_variable_id'] + variable_mention_mask = tensor_dict["variable_mention_mask"] + variable_mention_to_variable_id = tensor_dict["variable_mention_to_variable_id"] # (batch_size, max_variable_num) - variable_encoding_mask = tensor_dict['variable_encoding_mask'] - variable_mention_num = tensor_dict['variable_mention_num'] + variable_encoding_mask = tensor_dict["variable_encoding_mask"] + variable_mention_num = tensor_dict["variable_mention_num"] # # (batch_size, max_variable_mention_num, encoding_size) # variable_mention_encoding = torch.gather(code_token_encoding, 1, variable_mention_positions.unsqueeze(-1).expand(-1, -1, code_token_encoding.size(-1))) * variable_mention_positions_mask @@ -68,19 +80,27 @@ def forward(self, tensor_dict: Dict[str, torch.Tensor]): variable_num = variable_mention_num.size(1) encoding_size = code_token_encoding.size(-1) - variable_mention_encoding = code_token_encoding * variable_mention_mask.unsqueeze(-1) - variable_encoding = torch.zeros(tensor_dict['batch_size'], variable_num, encoding_size, device=self.device) - variable_encoding.scatter_add_(1, - variable_mention_to_variable_id.unsqueeze(-1).expand(-1, -1, encoding_size), - variable_mention_encoding) * variable_encoding_mask.unsqueeze(-1) - variable_encoding = variable_encoding / (variable_mention_num + (1. - variable_encoding_mask) * nn_util.SMALL_NUMBER).unsqueeze(-1) + variable_mention_encoding = ( + code_token_encoding * variable_mention_mask.unsqueeze(-1) + ) + variable_encoding = torch.zeros( + tensor_dict["batch_size"], variable_num, encoding_size, device=self.device + ) + variable_encoding.scatter_add_( + 1, + variable_mention_to_variable_id.unsqueeze(-1).expand(-1, -1, encoding_size), + variable_mention_encoding, + ) * variable_encoding_mask.unsqueeze(-1) + variable_encoding = variable_encoding / ( + variable_mention_num + (1.0 - variable_encoding_mask) * nn_util.SMALL_NUMBER + ).unsqueeze(-1) context_encoding = dict( variable_encoding=variable_encoding, code_token_encoding=code_token_encoding, code_token_mask=code_token_mask, last_states=last_states, - last_cells=last_cells + last_cells=last_cells, ) context_encoding.update(tensor_dict) @@ -99,12 +119,20 @@ def encode_sequence(self, code_sequence): # (batch_size) code_sequence_length = code_token_mask.sum(dim=-1).long() - sorted_seqs, sorted_seq_lens, restoration_indices, sorting_indices = nn_util.sort_batch_by_length(code_token_embedding, - code_sequence_length) + ( + sorted_seqs, + sorted_seq_lens, + restoration_indices, + sorting_indices, + ) = nn_util.sort_batch_by_length(code_token_embedding, code_sequence_length) - packed_question_embedding = pack_padded_sequence(sorted_seqs, sorted_seq_lens.data.tolist(), batch_first=True) + packed_question_embedding = pack_padded_sequence( + sorted_seqs, sorted_seq_lens.data.tolist(), batch_first=True + ) - sorted_encodings, (last_states, last_cells) = self.lstm_encoder(packed_question_embedding) + sorted_encodings, (last_states, last_cells) = self.lstm_encoder( + packed_question_embedding + ) sorted_encodings, _ = pad_packed_sequence(sorted_encodings, batch_first=True) # apply dropout to the last layer @@ -112,12 +140,18 @@ def encode_sequence(self, code_sequence): sorted_encodings = self.dropout(sorted_encodings) # (batch_size, question_len, hidden_size * 2) - restored_encodings = sorted_encodings.index_select(dim=0, index=restoration_indices) + restored_encodings = sorted_encodings.index_select( + dim=0, index=restoration_indices + ) # (num_layers, direction_num, batch_size, hidden_size) - last_states = last_states.view(self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size) + last_states = last_states.view( + self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size + ) last_states = last_states.index_select(dim=2, index=restoration_indices) - last_cells = last_cells.view(self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size) + last_cells = last_cells.view( + self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size + ) last_cells = last_cells.index_select(dim=2, index=restoration_indices) return restored_encodings, code_token_mask, (last_states, last_cells) @@ -127,25 +161,31 @@ def to_tensor_dict(cls, examples: List[Example]) -> Dict[str, torch.Tensor]: max_time_step = max(e.source_seq_length for e in examples) input = np.zeros((len(examples), max_time_step), dtype=np.int64) - variable_mention_to_variable_id = torch.zeros(len(examples), max_time_step, dtype=torch.long) + variable_mention_to_variable_id = torch.zeros( + len(examples), max_time_step, dtype=torch.long + ) variable_mention_mask = torch.zeros(len(examples), max_time_step) - variable_mention_num = torch.zeros(len(examples), max(len(e.ast.variables) for e in examples)) + variable_mention_num = torch.zeros( + len(examples), max(len(e.ast.variables) for e in examples) + ) variable_encoding_mask = torch.zeros(variable_mention_num.size()) for e_id, example in enumerate(examples): sub_tokens = example.sub_tokens - input[e_id, :len(sub_tokens)] = example.sub_token_ids + input[e_id, : len(sub_tokens)] = example.sub_token_ids variable_position_map = dict() var_name_to_id = {name: i for i, name in enumerate(example.ast.variables)} for i, sub_token in enumerate(sub_tokens): - if sub_token.startswith('@@') and sub_token.endswith('@@'): - old_var_name = sub_token[2: -2] - if old_var_name in var_name_to_id: # sometimes there are strings like `@@@@` + if sub_token.startswith("@@") and sub_token.endswith("@@"): + old_var_name = sub_token[2:-2] + if ( + old_var_name in var_name_to_id + ): # sometimes there are strings like `@@@@` var_id = var_name_to_id[old_var_name] variable_mention_to_variable_id[e_id, i] = var_id - variable_mention_mask[e_id, i] = 1. + variable_mention_mask[e_id, i] = 1.0 variable_position_map.setdefault(old_var_name, []).append(i) for var_id, var_name in enumerate(example.ast.variables): @@ -154,30 +194,38 @@ def to_tensor_dict(cls, examples: List[Example]) -> Dict[str, torch.Tensor]: variable_mention_num[e_id, var_id] = len(var_pos) except KeyError: variable_mention_num[e_id, var_id] = 1 - print(example.binary_file, f'variable [{var_name}] not found', file=sys.stderr) - - variable_encoding_mask[e_id, :len(example.ast.variables)] = 1. - - return dict(src_code_tokens=torch.from_numpy(input), - variable_mention_to_variable_id=variable_mention_to_variable_id, - variable_mention_mask=variable_mention_mask, - variable_mention_num=variable_mention_num, - variable_encoding_mask=variable_encoding_mask, - batch_size=len(examples)) + print( + example.binary_file, + f"variable [{var_name}] not found", + file=sys.stderr, + ) + + variable_encoding_mask[e_id, : len(example.ast.variables)] = 1.0 + + return dict( + src_code_tokens=torch.from_numpy(input), + variable_mention_to_variable_id=variable_mention_to_variable_id, + variable_mention_mask=variable_mention_mask, + variable_mention_num=variable_mention_num, + variable_encoding_mask=variable_encoding_mask, + batch_size=len(examples), + ) def get_decoder_init_state(self, context_encoder, config=None): - fwd_last_layer_cell = context_encoder['last_cells'][-1, 0] - bak_last_layer_cell = context_encoder['last_cells'][-1, 1] + fwd_last_layer_cell = context_encoder["last_cells"][-1, 0] + bak_last_layer_cell = context_encoder["last_cells"][-1, 1] - dec_init_cell = self.decoder_cell_init(torch.cat([fwd_last_layer_cell, bak_last_layer_cell], dim=-1)) + dec_init_cell = self.decoder_cell_init( + torch.cat([fwd_last_layer_cell, bak_last_layer_cell], dim=-1) + ) dec_init_state = torch.tanh(dec_init_cell) return dec_init_state, dec_init_cell - def get_attention_memory(self, context_encoding, att_target='terminal_nodes'): - assert att_target == 'terminal_nodes' + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + assert att_target == "terminal_nodes" - memory = context_encoding['code_token_encoding'] - mask = context_encoding['code_token_mask'] + memory = context_encoding["code_token_encoding"] + mask = context_encoding["code_token_mask"] return memory, mask diff --git a/neural-model/model/simple_decoder.py b/neural-model/model/simple_decoder.py index 375bdc6..127ecaa 100644 --- a/neural-model/model/simple_decoder.py +++ b/neural-model/model/simple_decoder.py @@ -1,12 +1,11 @@ from typing import Dict, List import torch -from torch import nn as nn - from model.decoder import Decoder -from utils import util, nn_util +from torch import nn as nn +from utils import nn_util, util from utils.ast import AbstractSyntaxTree -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab class SimpleDecoder(Decoder): @@ -14,32 +13,35 @@ def __init__(self, ast_node_encoding_size: int, vocab: Vocab): super(SimpleDecoder, self).__init__() self.vocab = vocab - self.state2names = nn.Linear(ast_node_encoding_size, len(vocab.target), bias=True) + self.state2names = nn.Linear( + ast_node_encoding_size, len(vocab.target), bias=True + ) self.config: Dict = None @classmethod def default_params(cls): - return { - 'vocab_file': None, - 'ast_node_encoding_size': 128 - } + return {"vocab_file": None, "ast_node_encoding_size": 128} @classmethod def build(cls, config): params = util.update(cls.default_params(), config) - vocab = torch.load(params['vocab_file']) - model = cls(params['ast_node_encoding_size'], vocab) + vocab = torch.load(params["vocab_file"]) + model = cls(params["ast_node_encoding_size"], vocab) model.config = params return model - def forward(self, src_ast_encoding: Dict[str, torch.Tensor], prediction_target: Dict[str, torch.Tensor]): + def forward( + self, + src_ast_encoding: Dict[str, torch.Tensor], + prediction_target: Dict[str, torch.Tensor], + ): """ Given a batch of encoded ASTs, compute the log-likelihood of generating all possible renamings """ # (all_var_node_num, tgt_vocab_size) - logits = self.state2names(src_ast_encoding['variable_master_node_encoding']) + logits = self.state2names(src_ast_encoding["variable_master_node_encoding"]) batched_p_names = torch.log_softmax(logits, dim=-1) # logits = self.state2names(src_ast_encoding) # p = torch.log_softmax(logits, dim=-1) @@ -49,33 +51,47 @@ def forward(self, src_ast_encoding: Dict[str, torch.Tensor], prediction_target: # batched_p_names.unsqueeze(-1).expand_as(src_ast_encoding.batch_size, -1, -1).scatter(idx, dim=1) # (batch_var_node_num) - packed_tgt_var_node_name_log_probs = torch.gather(batched_p_names, - dim=-1, - index=prediction_target['variable_tgt_name_id'].unsqueeze(-1)).squeeze(-1) + packed_tgt_var_node_name_log_probs = torch.gather( + batched_p_names, + dim=-1, + index=prediction_target["variable_tgt_name_id"].unsqueeze(-1), + ).squeeze(-1) # result = dict(context_encoding=context_encoding) result = dict() with torch.no_grad(): # compute ppl over renamed variables - renamed_var_mask = prediction_target['var_with_new_name_mask'] - unchanged_var_mask = prediction_target['auxiliary_var_mask'] - renamed_var_avg_ll = (packed_tgt_var_node_name_log_probs * renamed_var_mask).sum( - -1) / renamed_var_mask.sum() - unchanged_var_avg_ll = (packed_tgt_var_node_name_log_probs * unchanged_var_mask).sum( - -1) / unchanged_var_mask.sum() - result['rename_ppl'] = torch.exp(-renamed_var_avg_ll).item() - result['unchange_ppl'] = torch.exp(-unchanged_var_avg_ll).item() - - packed_tgt_var_node_name_log_probs = packed_tgt_var_node_name_log_probs * prediction_target['variable_tgt_name_weight'] + renamed_var_mask = prediction_target["var_with_new_name_mask"] + unchanged_var_mask = prediction_target["auxiliary_var_mask"] + renamed_var_avg_ll = ( + packed_tgt_var_node_name_log_probs * renamed_var_mask + ).sum(-1) / renamed_var_mask.sum() + unchanged_var_avg_ll = ( + packed_tgt_var_node_name_log_probs * unchanged_var_mask + ).sum(-1) / unchanged_var_mask.sum() + result["rename_ppl"] = torch.exp(-renamed_var_avg_ll).item() + result["unchange_ppl"] = torch.exp(-unchanged_var_avg_ll).item() + + packed_tgt_var_node_name_log_probs = ( + packed_tgt_var_node_name_log_probs + * prediction_target["variable_tgt_name_weight"] + ) # (batch_size, max_variable_node_num) - tgt_name_log_probs = packed_tgt_var_node_name_log_probs[src_ast_encoding['variable_encoding_restoration_indices']] - tgt_name_log_probs = tgt_name_log_probs * src_ast_encoding['variable_encoding_restoration_indices_mask'] + tgt_name_log_probs = packed_tgt_var_node_name_log_probs[ + src_ast_encoding["variable_encoding_restoration_indices"] + ] + tgt_name_log_probs = ( + tgt_name_log_probs + * src_ast_encoding["variable_encoding_restoration_indices_mask"] + ) # (batch_size) - ast_log_probs = tgt_name_log_probs.sum(dim=-1) / src_ast_encoding['variable_encoding_restoration_indices_mask'].sum(-1) + ast_log_probs = tgt_name_log_probs.sum(dim=-1) / src_ast_encoding[ + "variable_encoding_restoration_indices_mask" + ].sum(-1) - result['batch_log_prob'] = ast_log_probs + result["batch_log_prob"] = ast_log_probs return result @@ -89,7 +105,9 @@ def predict(self, source_asts: List[AbstractSyntaxTree]): context_encoding = self.encoder(tensor_dict) # (prediction_size, tgt_vocab_size) packed_var_name_log_probs = self.decoder(context_encoding) - best_var_name_log_probs, best_var_name_ids = torch.max(packed_var_name_log_probs, dim=-1) + best_var_name_log_probs, best_var_name_ids = torch.max( + packed_var_name_log_probs, dim=-1 + ) variable_rename_results = [] pred_node_ptr = 0 @@ -103,8 +121,10 @@ def predict(self, source_asts: List[AbstractSyntaxTree]): if new_var_name == SAME_VARIABLE_TOKEN: new_var_name = var_name - variable_rename_result[var_name] = {'new_name': new_var_name, - 'prob': var_name_prob} + variable_rename_result[var_name] = { + "new_name": new_var_name, + "prob": var_name_prob, + } pred_node_ptr += 1 @@ -112,4 +132,4 @@ def predict(self, source_asts: List[AbstractSyntaxTree]): assert pred_node_ptr == packed_var_name_log_probs.size(0) - return variable_rename_results \ No newline at end of file + return variable_rename_results diff --git a/neural-model/utils/ast.py b/neural-model/utils/ast.py index 21c5e70..766789f 100644 --- a/neural-model/utils/ast.py +++ b/neural-model/utils/ast.py @@ -1,18 +1,23 @@ from collections import OrderedDict from io import StringIO -import ujson as json from typing import Dict, List -import numpy as np - -import torch +import ujson as json from utils.util import cached_property from utils.vocab import VocabEntry class SyntaxNode(object): """represent a node on an AST""" - def __init__(self, node_id, node_type, address=None, children: List=None, named_fields: Dict=None): + + def __init__( + self, + node_id, + node_type, + address=None, + children: List = None, + named_fields: Dict = None, + ): self.node_id = node_id self.node_type = node_type self.address = address @@ -29,29 +34,35 @@ def __init__(self, node_id, node_type, address=None, children: List=None, named_ for child in children: self.add_child(child) - def add_child(self, child: 'SyntaxNode') -> None: + def add_child(self, child: "SyntaxNode") -> None: self.children.append(child) child.parent = self @classmethod - def from_json_dict(cls, json_dict: Dict) -> 'SyntaxNode': - named_fields = {k: v for k, v in json_dict.items() if k not in {'node_id', 'node_type', 'children', 'address', 'x', 'y', 'z'}} - - if 'x' in json_dict: - named_fields['x'] = SyntaxNode.from_json_dict(json_dict['x']) - if 'y' in json_dict: - named_fields['y'] = SyntaxNode.from_json_dict(json_dict['y']) - if 'z' in json_dict: - named_fields['z'] = SyntaxNode.from_json_dict(json_dict['z']) - - node = cls(json_dict['node_id'], - json_dict['node_type'], - json_dict['address'], - named_fields=named_fields) + def from_json_dict(cls, json_dict: Dict) -> "SyntaxNode": + named_fields = { + k: v + for k, v in json_dict.items() + if k not in {"node_id", "node_type", "children", "address", "x", "y", "z"} + } + + if "x" in json_dict: + named_fields["x"] = SyntaxNode.from_json_dict(json_dict["x"]) + if "y" in json_dict: + named_fields["y"] = SyntaxNode.from_json_dict(json_dict["y"]) + if "z" in json_dict: + named_fields["z"] = SyntaxNode.from_json_dict(json_dict["z"]) + + node = cls( + json_dict["node_id"], + json_dict["node_type"], + json_dict["address"], + named_fields=named_fields, + ) children_list = [] - if 'children' in json_dict: - children_list.extend(json_dict['children']) + if "children" in json_dict: + children_list.extend(json_dict["children"]) for child_dict in children_list: child_node = SyntaxNode.from_json_dict(child_dict) @@ -60,13 +71,13 @@ def from_json_dict(cls, json_dict: Dict) -> 'SyntaxNode': return node def to_json_dict(self): - json_dict = dict(node_id=self.node_id, - node_type=self.node_type, - address=self.address) + json_dict = dict( + node_id=self.node_id, node_type=self.node_type, address=self.address + ) for named_filed in self.named_fields: val = getattr(self, named_filed) - if named_filed in ('x', 'y', 'z'): + if named_filed in ("x", "y", "z"): json_dict[named_filed] = val.to_json_dict() else: json_dict[named_filed] = val @@ -75,17 +86,22 @@ def to_json_dict(self): children = [] for child in self.children: children.append(child.to_json_dict()) - json_dict['children'] = children + json_dict["children"] = children return json_dict @property def is_variable_node(self): - return self.node_type == 'var' + return self.node_type == "var" @property def is_terminal_node(self): - return not hasattr(self, 'x') and not hasattr(self, 'y') and not hasattr(self, 'z') and not self.children + return ( + not hasattr(self, "x") + and not hasattr(self, "y") + and not hasattr(self, "z") + and not self.children + ) @cached_property def size(self): @@ -97,11 +113,11 @@ def size(self): @property def member_nodes(self): - if hasattr(self, 'x'): + if hasattr(self, "x"): yield self.x - if hasattr(self, 'y'): + if hasattr(self, "y"): yield self.y - if hasattr(self, 'z'): + if hasattr(self, "z"): yield self.z for child in self.children: @@ -109,14 +125,14 @@ def member_nodes(self): @property def named_succeeding_fields(self): - if hasattr(self, 'x'): - yield 'x', self.x - if hasattr(self, 'y'): - yield 'y', self.y - if hasattr(self, 'z'): - yield 'z', self.z + if hasattr(self, "x"): + yield "x", self.x + if hasattr(self, "y"): + yield "y", self.y + if hasattr(self, "z"): + yield "z", self.z - yield 'children', self.children + yield "children", self.children @property def descendant_nodes(self): @@ -125,6 +141,7 @@ def _visit(node): for member_node in node.member_nodes: yield from _visit(member_node) + yield from _visit(self) def __iter__(self): @@ -138,19 +155,20 @@ def __hash__(self): return code def __eq__(self, other): - if not isinstance(other, self.__class__): return False - - if self.node_type != other.node_type: return False - if self.address != other.address: return False - if self.node_id != other.node_id: return False - if len(self.children) != len(other.children): return False + if ( + not isinstance(other, self.__class__) + or self.node_type != other.node_type + or self.address != other.address + or self.node_id != other.node_id + or self.named_fields != other.named_fields + or len(self.children) != len(other.children) + ): + return False for i in range(len(self.children)): if self.children[i] != other.children[i]: return False - if self.named_fields != other.named_fields: return False - return True def __ne__(self, other): @@ -162,44 +180,51 @@ def to_string(self, sb=None): is_root = True sb = StringIO() - sb.write(f'(Node{self.node_id}-{self.node_type}') + sb.write(f"(Node{self.node_id}-{self.node_type}") for key in self.named_fields: val = getattr(self, key) - if key not in ('x', 'y', 'z'): - sb.write('-') - sb.write(f'{key}:{val}'.replace(' ', '_').replace('(', '_').replace(')', '_')) + if key not in ("x", "y", "z"): + sb.write("-") + sb.write( + f"{key}:{val}".replace(" ", "_").replace("(", "_").replace(")", "_") + ) # sb.write(f'Node-{self.node_id}-{self.node_type}') for field_name, node in self.named_succeeding_fields: - if field_name in ('x', 'y'): - sb.write(f' ({field_name} ') + if field_name in ("x", "y"): + sb.write(f" ({field_name} ") node.to_string(sb) - sb.write(')') # of x + sb.write(")") # of x elif self.children: - sb.write(' (children ') + sb.write(" (children ") for child in self.children: - sb.write(' ') + sb.write(" ") child.to_string(sb) - sb.write(')') # of children field + sb.write(")") # of children field - sb.write(')') # of node + sb.write(")") # of node if is_root: return sb.getvalue() def __str__(self): - return f'Node {self.node_id} {self.node_type}@{self.address}' + return f"Node {self.node_id} {self.node_type}@{self.address}" __repr__ = __str__ class TerminalNode(SyntaxNode): - """a terminal AST node representing variables or other terminal syntax tokens""" + """a terminal AST node representing variables or other terminal syntax + tokens + """ + pass class AbstractSyntaxTree(object): - def __init__(self, root: SyntaxNode, compilation_unit: str = None, code: str = None): + def __init__( + self, root: SyntaxNode, compilation_unit: str = None, code: str = None + ): self.root = root self.compilation_unit = compilation_unit self.code = code @@ -213,10 +238,14 @@ def __init__(self, root: SyntaxNode, compilation_unit: str = None, code: str = N self._init_index() @classmethod - def from_json_dict(cls, json_dict: Dict) -> 'AbstractSyntaxTree': - root = SyntaxNode.from_json_dict(json_dict['ast']) - root.name = json_dict['function'] - tree = cls(root, compilation_unit=json_dict['function'], code=json_dict['raw_code'] if 'raw_code' in json_dict else None) + def from_json_dict(cls, json_dict: Dict) -> "AbstractSyntaxTree": + root = SyntaxNode.from_json_dict(json_dict["ast"]) + root.name = json_dict["function"] + tree = cls( + root, + compilation_unit=json_dict["function"], + code=json_dict.get("raw_code", None), + ) return tree @@ -248,19 +277,21 @@ def _index_sub_tree(node: SyntaxNode, parent_node: SyntaxNode): _index_sub_tree(self.root, None) - setattr(self, 'adjacency_list', adj_list) - setattr(self, 'id_to_node', id2node) - setattr(self, 'adjacent_terminal_nodes', []) # TODO: implement this! - setattr(self, 'variable_nodes', variable_nodes) - setattr(self, 'variables', variables) - terminal_nodes.sort(key=lambda n: n.node_id) # TODO: change to address based! - setattr(self, 'terminal_nodes', terminal_nodes) + setattr(self, "adjacency_list", adj_list) + setattr(self, "id_to_node", id2node) + # TODO: implement this! + setattr(self, "adjacent_terminal_nodes", []) + setattr(self, "variable_nodes", variable_nodes) + setattr(self, "variables", variables) + # TODO: change to address based! + terminal_nodes.sort(key=lambda n: n.node_id) + setattr(self, "terminal_nodes", terminal_nodes) def __iter__(self): return iter((node for node in self.id_to_node.values())) -if __name__ == '__main__': +if __name__ == "__main__": json_str = """ { "function": "ft_strncat", @@ -693,7 +724,7 @@ def __iter__(self): } """ json_dict = json.loads(json_str) - tree = SyntaxNode.from_json_dict(json_dict['root']) + tree = SyntaxNode.from_json_dict(json_dict["root"]) tree_reconstr = SyntaxNode.from_json_dict(tree.to_json_dict()) assert tree_reconstr == tree @@ -703,4 +734,4 @@ def __iter__(self): from utils.code_processing import annotate_type annotate_type(tree_reconstr) - print(tree_reconstr.to_json_dict()) \ No newline at end of file + print(tree_reconstr.to_json_dict()) diff --git a/neural-model/utils/code_processing.py b/neural-model/utils/code_processing.py index 43d1c7b..0cf2775 100644 --- a/neural-model/utils/code_processing.py +++ b/neural-model/utils/code_processing.py @@ -1,30 +1,29 @@ import re -from typing import List, Set +from typing import Set from utils.ast import SyntaxNode -from utils.lexer import * +from utils.lexer import Lexer, Token - -VARIABLE_ANNOTATION = re.compile(r'@@\w+@@(\w+)@@\w+') +VARIABLE_ANNOTATION = re.compile(r"@@\w+@@(\w+)@@\w+") def canonicalize_code(code): - code = re.sub('//.*?\\n|/\\*.*?\\*/', '\\n', code, flags=re.S) - lines = [l.rstrip() for l in code.split('\\n')] - code = '\\n'.join(lines) - code = re.sub('@@\\w+@@(\\w+)@@\\w+', '\\g<1>', code) + code = re.sub("//.*?\\n|/\\*.*?\\*/", "\\n", code, flags=re.S) + lines = [l.rstrip() for l in code.split("\\n")] + code = "\\n".join(lines) + code = re.sub("@@\\w+@@(\\w+)@@\\w+", "\\g<1>", code) return code def canonicalize_constants(root: SyntaxNode) -> None: def _visit(node): - if node.node_type == 'obj' and node.type == 'char *': - node.name = 'STRING' - elif node.node_type == 'num': - node.name = 'NUMBER' - elif node.node_type == 'fnum': - node.name = 'FLOAT' + if node.node_type == "obj" and node.type == "char *": + node.name = "STRING" + elif node.node_type == "num": + node.name = "NUMBER" + elif node.node_type == "fnum": + node.name = "FLOAT" for child in node.member_nodes: _visit(child) @@ -34,11 +33,11 @@ def _visit(node): def annotate_type(root: SyntaxNode) -> None: def _visit(node): - if hasattr(node, 'type'): - type_tokens = [t[1].lstrip('_') for t in Lexer(node.type).get_tokens()] - type_tokens = [t for t in type_tokens if t not in ('(', ')')] - node.named_fields.add('type_tokens') - setattr(node, 'type_tokens', type_tokens) + if hasattr(node, "type"): + type_tokens = [t[1].lstrip("_") for t in Lexer(node.type).get_tokens()] + type_tokens = [t for t in type_tokens if t not in ("(", ")")] + node.named_fields.add("type_tokens") + setattr(node, "type_tokens", type_tokens) for child in node.member_nodes: _visit(child) @@ -49,35 +48,37 @@ def _visit(node): VAR_ID_REGEX = re.compile(r"@@(VAR_\d+)@@") -def preprocess_ast(root: SyntaxNode, preprocessors: Set[str] = None, code: str = None) -> None: +def preprocess_ast( + root: SyntaxNode, preprocessors: Set[str] = None, code: str = None +) -> None: if preprocessors is None: - preprocessors = {'annotate_type', 'canonicalize_constant', 'annotate_arg'} + preprocessors = {"annotate_type", "canonicalize_constant", "annotate_arg"} arg_var_ids = None - if 'annotate_arg' in preprocessors: - first_line = code[:code.index('\n')] + if "annotate_arg" in preprocessors: + first_line = code[: code.index("\n")] arg_var_ids = set(VAR_ID_REGEX.findall(first_line)) def _visit(node): - if 'annotate_type' in preprocessors: - if node.node_type == 'obj' and node.type == 'char *': - node.name = 'STRING' - elif node.node_type == 'num': - node.name = 'NUMBER' - elif node.node_type == 'fnum': - node.name = 'FLOAT' - - if 'canonicalize_constant' in preprocessors: - if hasattr(node, 'type'): - type_tokens = [t[1].lstrip('_') for t in Lexer(node.type).get_tokens()] - type_tokens = [t for t in type_tokens if t not in ('(', ')')] - node.named_fields.add('type_tokens') - setattr(node, 'type_tokens', type_tokens) - - if 'annotate_arg' in preprocessors: - if node.node_type == 'var': - node.named_fields.add('is_arg') - setattr(node, 'is_arg', node.var_id in arg_var_ids) + if "annotate_type" in preprocessors: + if node.node_type == "obj" and node.type == "char *": + node.name = "STRING" + elif node.node_type == "num": + node.name = "NUMBER" + elif node.node_type == "fnum": + node.name = "FLOAT" + + if "canonicalize_constant" in preprocessors: + if hasattr(node, "type"): + type_tokens = [t[1].lstrip("_") for t in Lexer(node.type).get_tokens()] + type_tokens = [t for t in type_tokens if t not in ("(", ")")] + node.named_fields.add("type_tokens") + setattr(node, "type_tokens", type_tokens) + + if "annotate_arg" in preprocessors: + if node.node_type == "var": + node.named_fields.add("is_arg") + setattr(node, "is_arg", node.var_id in arg_var_ids) for child in node.member_nodes: _visit(child) @@ -90,14 +91,13 @@ def tokenize_raw_code(raw_code): tokens = [] for token_type, token in lexer.get_tokens(): if token_type in Token.Literal: - token = str(token_type).split('.')[2] + token = str(token_type).split(".")[2] if token_type == Token.Placeholder.Var: m = VARIABLE_ANNOTATION.match(token) old_name = m.group(1) - token = '@@' + old_name + '@@' + token = "@@" + old_name + "@@" tokens.append(token) return tokens - diff --git a/neural-model/utils/dataset.py b/neural-model/utils/dataset.py index 5aad665..ad535e7 100644 --- a/neural-model/utils/dataset.py +++ b/neural-model/utils/dataset.py @@ -1,36 +1,32 @@ +import gc import glob +import multiprocessing +import os import pickle +import queue +import random +import resource import sys -import os -import gc -import time -import ujson as json import tarfile -from typing import Iterable, List, Dict, Union, Tuple -import multiprocessing import threading -import queue +import time +from typing import Dict, Iterable, List, Tuple, Union -from tqdm import tqdm import numpy as np - +import sentencepiece as spm +import torch +import torch.multiprocessing as torch_mp +import ujson as json +from tqdm import tqdm from utils import nn_util from utils.ast import AbstractSyntaxTree, SyntaxNode from utils.code_processing import annotate_type from utils.graph import PackedGraph -from utils.vocab import VocabEntry, SAME_VARIABLE_TOKEN, Vocab -import sentencepiece as spm -import random - -import torch -import torch.multiprocessing as torch_mp - +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab, VocabEntry batcher_sync_msg = None -torch.multiprocessing.set_sharing_strategy('file_system') - -import resource +torch.multiprocessing.set_sharing_strategy("file_system") rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1])) @@ -51,14 +47,14 @@ def from_json_dict(cls, json_dict, **kwargs): for var_name, var_nodes in tree.variables.items(): variable_name_map[var_name] = var_nodes[0].new_name - if 'test_meta' in json_dict: - kwargs['test_meta'] = json_dict['test_meta'] + if "test_meta" in json_dict: + kwargs["test_meta"] = json_dict["test_meta"] return cls(tree, variable_name_map, **kwargs) class Batch(object): - __slots__ = ('examples', 'tensor_dict') + __slots__ = ("examples", "tensor_dict") def __init__(self, examples, tensor_dict): self.examples = examples @@ -66,7 +62,7 @@ def __init__(self, examples, tensor_dict): @property def size(self): - return self.tensor_dict['batch_size'] + return self.tensor_dict["batch_size"] class Batcher(object): @@ -75,34 +71,48 @@ def __init__(self, config, train=True): self.train = train # model specific config - self.is_ensemble = config['encoder']['type'] == 'EnsembleModel' + self.is_ensemble = config["encoder"]["type"] == "EnsembleModel" if not self.is_ensemble: - self.vocab = Vocab.load(config['data']['vocab_file']) + self.vocab = Vocab.load(config["data"]["vocab_file"]) self.grammar = self.vocab.grammar - self.use_seq_encoder = config['encoder']['type'] == 'SequentialEncoder' - self.use_hybrid_encoder = config['encoder']['type'] == 'HybridEncoder' - self.init_gnn_with_seq_encoding = config['encoder']['type'] == 'GraphASTEncoder' and config['encoder']['init_with_seq_encoding'] + self.use_seq_encoder = config["encoder"]["type"] == "SequentialEncoder" + self.use_hybrid_encoder = config["encoder"]["type"] == "HybridEncoder" + self.init_gnn_with_seq_encoding = ( + config["encoder"]["type"] == "GraphASTEncoder" + and config["encoder"]["init_with_seq_encoding"] + ) @property def annotate_sequential_input(self): - return self.use_seq_encoder or self.use_hybrid_encoder or self.init_gnn_with_seq_encoding + return ( + self.use_seq_encoder + or self.use_hybrid_encoder + or self.init_gnn_with_seq_encoding + ) def annotate_example(self, example) -> Example: - """annotate examples by populating specific fields, useful for sorting examples or batching""" - # for ensemble models, it will be annotated by the batcher for each specific class + """Annotate examples by populating specific fields, useful for sorting + examples or batching. + """ + # for ensemble models, it will be annotated by the batcher for each + # specific class if self.is_ensemble: return example if self.annotate_sequential_input: src_bpe_model = self.vocab.source_tokens.subtoken_model snippet = example.code_tokens - snippet = ' '.join(snippet) - sub_tokens = [''] + src_bpe_model.encode_as_pieces(snippet) + [''] - sub_token_ids = [src_bpe_model.bos_id()] + src_bpe_model.encode_as_ids(snippet) + [src_bpe_model.eos_id()] - setattr(example, 'sub_tokens', sub_tokens) - setattr(example, 'sub_token_ids', sub_token_ids) - setattr(example, 'source_seq_length', len(sub_tokens)) + snippet = " ".join(snippet) + sub_tokens = [""] + src_bpe_model.encode_as_pieces(snippet) + [""] + sub_token_ids = ( + [src_bpe_model.bos_id()] + + src_bpe_model.encode_as_ids(snippet) + + [src_bpe_model.eos_id()] + ) + setattr(example, "sub_tokens", sub_tokens) + setattr(example, "sub_token_ids", sub_token_ids) + setattr(example, "source_seq_length", len(sub_tokens)) tgt_bpe_model = self.vocab.target.subtoken_model eov_id = tgt_bpe_model.eos_id() @@ -116,8 +126,8 @@ def annotate_example(self, example) -> Example: variable_name_subtoken_map[old_name] = subtoken_ids tgt_pred_seq_len += len(subtoken_ids) - setattr(example, 'variable_name_subtoken_map', variable_name_subtoken_map) - setattr(example, 'target_prediction_seq_length', tgt_pred_seq_len) + setattr(example, "variable_name_subtoken_map", variable_name_subtoken_map) + setattr(example, "target_prediction_seq_length", tgt_pred_seq_len) return example @@ -143,57 +153,73 @@ def get_batch_size(self, examples: List[Example]): else: return len(examples) * max(e.target_prediction_seq_length for e in examples) - def to_tensor_dict(self, examples: List[Example], return_prediction_target=True) -> Dict[str, torch.Tensor]: - from model.sequential_encoder import SequentialEncoder + def to_tensor_dict( + self, examples: List[Example], return_prediction_target=True + ) -> Dict[str, torch.Tensor]: from model.graph_encoder import GraphASTEncoder + from model.sequential_encoder import SequentialEncoder - if not hasattr(examples[0], 'target_prediction_seq_length'): + if not hasattr(examples[0], "target_prediction_seq_length"): for example in examples: self.annotate_example(example) - if self.config['encoder']['type'] == 'GraphASTEncoder': - init_with_seq_encoding = self.config['encoder']['init_with_seq_encoding'] - packed_graph, tensor_dict = GraphASTEncoder.to_packed_graph([e.ast for e in examples], - connections=self.config['encoder']['connections'], - init_with_seq_encoding=init_with_seq_encoding) + if self.config["encoder"]["type"] == "GraphASTEncoder": + init_with_seq_encoding = self.config["encoder"]["init_with_seq_encoding"] + packed_graph, tensor_dict = GraphASTEncoder.to_packed_graph( + [e.ast for e in examples], + connections=self.config["encoder"]["connections"], + init_with_seq_encoding=init_with_seq_encoding, + ) if init_with_seq_encoding: seq_tensor_dict = SequentialEncoder.to_tensor_dict(examples) - tensor_dict['seq_encoder_input'] = seq_tensor_dict + tensor_dict["seq_encoder_input"] = seq_tensor_dict - _tensors = GraphASTEncoder.to_tensor_dict(packed_graph, - self.grammar, self.vocab) + _tensors = GraphASTEncoder.to_tensor_dict( + packed_graph, self.grammar, self.vocab + ) tensor_dict.update(_tensors) - elif self.config['encoder']['type'] == 'SequentialEncoder': + elif self.config["encoder"]["type"] == "SequentialEncoder": tensor_dict = SequentialEncoder.to_tensor_dict(examples) - elif self.config['encoder']['type'] == 'HybridEncoder': - packed_graph, gnn_tensor_dict = GraphASTEncoder.to_packed_graph([e.ast for e in examples], - connections=self.config['encoder']['graph_encoder']['connections']) - gnn_tensors = GraphASTEncoder.to_tensor_dict(packed_graph, self.grammar, self.vocab) + elif self.config["encoder"]["type"] == "HybridEncoder": + packed_graph, gnn_tensor_dict = GraphASTEncoder.to_packed_graph( + [e.ast for e in examples], + connections=self.config["encoder"]["graph_encoder"]["connections"], + ) + gnn_tensors = GraphASTEncoder.to_tensor_dict( + packed_graph, self.grammar, self.vocab + ) gnn_tensor_dict.update(gnn_tensors) seq_tensor_dict = SequentialEncoder.to_tensor_dict(examples) - tensor_dict = {'graph_encoder_input': gnn_tensor_dict, - 'seq_encoder_input': seq_tensor_dict} + tensor_dict = { + "graph_encoder_input": gnn_tensor_dict, + "seq_encoder_input": seq_tensor_dict, + } else: - raise ValueError('UnknownEncoderType') + raise ValueError("UnknownEncoderType") if self.train or return_prediction_target: prediction_target = self.to_batched_prediction_target(examples) - tensor_dict['prediction_target'] = prediction_target + tensor_dict["prediction_target"] = prediction_target if not self.train: - if hasattr(examples[0], 'test_meta'): - tensor_dict['test_meta'] = [e.test_meta for e in examples] + if hasattr(examples[0], "test_meta"): + tensor_dict["test_meta"] = [e.test_meta for e in examples] - tensor_dict['batch_size'] = len(examples) + tensor_dict["batch_size"] = len(examples) num_elements = nn_util.get_tensor_dict_size(tensor_dict) - tensor_dict['num_elements'] = num_elements + tensor_dict["num_elements"] = num_elements return tensor_dict - def to_batch(self, examples: List[Example], return_examples=False, return_prediction_target=True) -> Batch: + def to_batch( + self, + examples: List[Example], + return_examples=False, + return_prediction_target=True, + ) -> Batch: if self.is_ensemble: # do not perform tensorization for the parent ensemble model tensor_dict = None @@ -211,22 +237,13 @@ def to_batch(self, examples: List[Example], return_examples=False, return_predic def to_batched_prediction_target(self, examples: List[Example]): batch_size = len(examples) - unchanged_var_weight = self.config['train']['unchanged_variable_weight'] + unchanged_var_weight = self.config["train"]["unchanged_variable_weight"] use_bpe_for_var_name = self.vocab.target.subtoken_model is not None variable_name_subtoken_maps = [] if use_bpe_for_var_name: - # eov_id = self.vocab.target.subtoken_model.eos_id() - # for var_name_map in variable_name_maps: - # var_name_subtoken_map = dict() - # for old_name, new_name in var_name_map.items(): - # if old_name == new_name: - # subtoken_ids = [self.vocab.target[SAME_VARIABLE_TOKEN], eov_id] - # else: - # subtoken_ids = self.vocab.target.subtoken_model.encode_as_ids(new_name) + [eov_id] - # var_name_subtoken_map[old_name] = subtoken_ids - variable_name_subtoken_maps = [e.variable_name_subtoken_map for e in examples] + var_name_subtoken_maps = [e.variable_name_subtoken_map for e in examples] else: for example in examples: var_name_map = example.variable_name_map @@ -239,53 +256,71 @@ def to_batched_prediction_target(self, examples: List[Example]): var_name_subtoken_map[old_name] = subtoken_ids variable_name_subtoken_maps.append(var_name_subtoken_map) - max_pred_timestep = max(sum(len(val) for val in x.values()) for x in variable_name_subtoken_maps) + max_pred_timestep = max( + sum(len(val) for val in x.values()) for x in variable_name_subtoken_maps + ) - target_variable_encoding_indices = torch.zeros(batch_size, max_pred_timestep, dtype=torch.long) - target_variable_encoding_indices_mask = torch.zeros(batch_size, max_pred_timestep) + target_variable_encoding_indices = torch.zeros( + batch_size, max_pred_timestep, dtype=torch.long + ) + target_variable_encoding_indices_mask = torch.zeros( + batch_size, max_pred_timestep + ) - variable_tgt_name_id = torch.zeros(batch_size, max_pred_timestep, dtype=torch.long) + variable_tgt_name_id = torch.zeros( + batch_size, max_pred_timestep, dtype=torch.long + ) variable_tgt_name_weight = torch.zeros(batch_size, max_pred_timestep) - var_with_new_name_mask = torch.zeros(batch_size, max_pred_timestep) + variable_with_new_name_mask = torch.zeros(batch_size, max_pred_timestep) auxiliary_var_mask = torch.zeros(batch_size, max_pred_timestep) variable_master_node_ptr = 0 for e_id, example in enumerate(examples): ast = example.ast var_name_map = example.variable_name_map - _var_node_ids = [] - _tgt_name_ids = [] + variable_ptr = 0 for var_id, var_name in enumerate(ast.variables): - new_var_name_subtoken_ids = variable_name_subtoken_maps[e_id][var_name] + new_variable_name_subtoken_ids = variable_name_subtoken_maps[e_id][var_name] variable_end_ptr = variable_ptr + len(new_var_name_subtoken_ids) - variable_tgt_name_id[e_id, variable_ptr: variable_end_ptr] = torch.tensor(new_var_name_subtoken_ids, dtype=torch.long) + variable_tgt_name_id[ + e_id, variable_ptr:variable_end_ptr + ] = torch.tensor(new_var_name_subtoken_ids, dtype=torch.long) if var_name == var_name_map[var_name]: - auxiliary_var_mask[e_id, variable_ptr: variable_end_ptr] = 1. - variable_tgt_name_weight[e_id, variable_ptr: variable_end_ptr] = unchanged_var_weight + auxiliary_var_mask[e_id, var_ptr:var_end_ptr] = 1.0 + + variable_tgt_name_weight[ + e_id, variable_ptr:variable_end_ptr + ] = unchanged_var_weight else: - var_with_new_name_mask[e_id, variable_ptr: variable_end_ptr] = 1. - variable_tgt_name_weight[e_id, variable_ptr: variable_end_ptr] = 1. + variable_with_new_name_mask[e_id, variable_ptr:variable_end_ptr] = 1.0 + variable_tgt_name_weight[e_id, variable_ptr:variable_end_ptr] = 1.0 - target_variable_encoding_indices[e_id, variable_ptr: variable_end_ptr] = var_id # variable_master_node_ptr + target_variable_encoding_indices[ + e_id, variable_ptr:variable_end_ptr + ] = var_id # variable_master_node_ptr variable_master_node_ptr += 1 variable_ptr = variable_end_ptr - target_variable_encoding_indices_mask[e_id, :variable_ptr] = 1. + target_variable_encoding_indices_mask[e_id, :variable_ptr] = 1.0 - return dict(variable_tgt_name_id=variable_tgt_name_id, - variable_tgt_name_weight=variable_tgt_name_weight, - var_with_new_name_mask=var_with_new_name_mask, - auxiliary_var_mask=auxiliary_var_mask, - target_variable_encoding_indices=target_variable_encoding_indices, - target_variable_encoding_indices_mask=target_variable_encoding_indices_mask) + return dict( + variable_tgt_name_id=variable_tgt_name_id, + variable_tgt_name_weight=variable_tgt_name_weight, + variable_with_new_name_mask=var_with_new_name_mask, + auxiliary_var_mask=auxiliary_var_mask, + target_variable_encoding_indices=target_variable_encoding_indices, + target_variable_encoding_indices_mask=target_variable_encoding_indices_mask, + ) -def get_json_iterator_from_tar_file(file_paths, shuffle=False, progress=False, group_by=None, buffer=True) -> Iterable: - assert group_by in (None, 'binary_file') +def get_json_iterator_from_tar_file( + file_paths, shuffle=False, progress=False, group_by=None, buffer=True +) -> Iterable: + assert group_by in (None, "binary_file") # if shuffle: # assert buffer is False @@ -298,13 +333,12 @@ def get_json_iterator_from_tar_file(file_paths, shuffle=False, progress=False, g for file_path in file_paths: payloads = [] t1 = time.time() - with tarfile.open(file_path, 'r') as f: - files = [x.name for x in f.getmembers() if x.name.endswith('.jsonl')] - # if shuffle: - # np.random.shuffle(files) - - if progress: file_iter = tqdm(files, file=sys.stdout) - else: file_iter = files + with tarfile.open(file_path, "r") as f: + files = [x.name for x in f.getmembers() if x.name.endswith(".jsonl")] + if progress: + file_iter = tqdm(files, file=sys.stdout) + else: + file_iter = files for filename in file_iter: jsonl_file = f.extractfile(filename) @@ -313,28 +347,40 @@ def get_json_iterator_from_tar_file(file_paths, shuffle=False, progress=False, g for line_no, tree_encoding_line in enumerate(jsonl_file): # if tree_encoding_line.decode().startswith('{'): # tree_json_dict = json.loads(tree_encoding_line) - payload = tree_encoding_line, dict(file_name=filename, line_num=line_no) + payload = ( + tree_encoding_line, + dict(file_name=filename, line_num=line_no), + ) if buffer: payloads.append(payload) else: yield payload - elif group_by == 'binary_file': - lines = [(l.decode().strip(), dict(file_name=filename, line_num=line_no)) - for line_no, l in enumerate(jsonl_file)] + elif group_by == "binary_file": + lines = [ + ( + l.decode().strip(), + dict(file_name=filename, line_num=line_no), + ) + for line_no, l in enumerate(jsonl_file) + ] yield lines if shuffle: np.random.shuffle(payloads) - print(f'load shard {file_path} took {time.time() - t1:.4f}s', file=sys.stderr) + print(f"load shard {file_path} took {time.time() - t1:.4f}s", file=sys.stderr) for payload in payloads: yield payload -def json_line_reader(file_path, queue, worker_num, shuffle, progress, group_by=None, buffer=True): - for json_str in get_json_iterator_from_tar_file(file_path, shuffle, progress, group_by=group_by, buffer=buffer): +def json_line_reader( + file_path, queue, worker_num, shuffle, progress, group_by=None, buffer=True +): + for json_str in get_json_iterator_from_tar_file( + file_path, shuffle, progress, group_by=group_by, buffer=buffer + ): queue.put(json_str) for i in range(worker_num): @@ -342,23 +388,30 @@ def json_line_reader(file_path, queue, worker_num, shuffle, progress, group_by=N def is_valid_training_example(example): - if hasattr(example, 'target_prediction_seq_length'): - if example.target_prediction_seq_length >= 200: return False + if hasattr(example, "target_prediction_seq_length"): + if example.target_prediction_seq_length >= 200: + return False - return example.ast.size < 300 and \ - len(example.variable_name_map) > 0 and \ - any(k != v for k, v in example.variable_name_map.items()) + return ( + example.ast.size < 300 + and len(example.variable_name_map) > 0 + and any(k != v for k, v in example.variable_name_map.items()) + ) def example_generator(json_queue, example_queue, consumer_num=1): while True: payload = json_queue.get() - if payload is None: break + if payload is None: + break json_str, meta = payload tree_json_dict = json.loads(json_str) - if 'code_tokens' in tree_json_dict: - example = Example.from_json_dict(tree_json_dict, binary_file=meta, code_tokens=tree_json_dict['code_tokens']) + if "code_tokens" in tree_json_dict: + code_tokens = tree_json_dict["code_tokens"] + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, code_tokens=code_tokens + ) else: example = Example.from_json_dict(tree_json_dict, binary_file=meta) @@ -370,12 +423,21 @@ def example_generator(json_queue, example_queue, consumer_num=1): # print('[Example Generator] example generator process quit!') -def example_to_batch(json_queue, batched_examples_queue, batch_size, train, config, worker_manager_lock, return_examples=False, return_prediction_target=True): +def example_to_batch( + json_queue, + batched_examples_queue, + batch_size, + train, + config, + worker_manager_lock, + return_examples=False, + return_prediction_target=True, +): batcher = Batcher(config, train) - buffer_size = config['train']['buffer_size'] + buffer_size = config["train"]["buffer_size"] buffer = [] - print(f'[ExampleToBatch] pid={os.getpid()}', file=sys.stderr) + print(f"[ExampleToBatch] pid={os.getpid()}", file=sys.stderr) def _generate_batches(): # buffer.sort(key=batcher.train_example_sort_key) @@ -399,21 +461,24 @@ def _generate_batches(): random.shuffle(batches) for batch_examples in batches: - batch = batcher.to_batch(batch_examples, return_examples=return_examples, return_prediction_target=return_prediction_target) + batch = batcher.to_batch( + batch_examples, + return_examples=return_examples, + return_prediction_target=return_prediction_target, + ) # while batched_examples_queue.qsize() > 100: # time.sleep(10) # print(batch.tensor_dict['num_elements']) while worker_manager_lock.value == 1: time.sleep(0.2) batched_examples_queue.put(batch) - # print(f'[ExampleToBatch] batched examples queue size {batched_examples_queue.qsize()}', file=sys.stderr) buffer.clear() gc.collect() finished = False while True: - t1 = time.time() + # t1 = time.time() while len(buffer) < buffer_size: payload = json_queue.get() @@ -424,9 +489,11 @@ def _generate_batches(): json_str, meta = payload tree_json_dict = json.loads(json_str) - if 'code_tokens' in tree_json_dict: - example = Example.from_json_dict(tree_json_dict, binary_file=meta, - code_tokens=tree_json_dict['code_tokens']) + if "code_tokens" in tree_json_dict: + code_tokens = tree_json_dict["code_tokens"] + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, code_tokens=code_tokens + ) else: example = Example.from_json_dict(tree_json_dict, binary_file=meta) @@ -450,33 +517,35 @@ def _generate_batches(): while batcher_sync_msg.value == 0: time.sleep(1) - print(f'[ExampleToBatch] quit', file=sys.stderr) + print(f"[ExampleToBatch] quit", file=sys.stderr) sys.stderr.flush() -def worker_manager(worker_result_queue, out_queue, num_workers, worker_manager_lock, buffer_size): +def worker_manager( + worker_result_queue, out_queue, num_workers, worker_manager_lock, buffer_size +): num_finished_workers = 0 patience = 0 prev_queue_size = -1 while True: finished = False - # t0 = time.time() try: queue_size = worker_result_queue.qsize() - except: - queue_size = 999999 # just trigger data loading, for max os X - # print(f'[LocalWorkerManager] queue size={queue_size}, patience={patience}', file=sys.stderr) - if (queue_size > buffer_size or patience >= 10) and out_queue.qsize() < buffer_size: + except Exception: + # just trigger data loading, for max os X + queue_size = 999999 + + if ( + queue_size > buffer_size or patience >= 10 + ) and out_queue.qsize() < buffer_size: worker_manager_lock.value = 1 patience = 0 - # print(f'[LocalWorkerManager] start loading {queue_size} batches...', file=sys.stderr) i = 0 while not worker_result_queue.empty() and i < buffer_size: batch = worker_result_queue.get() - # print(f'[LocalWorkerManager] {time.time() - t0} took to load a batch, size={worker_result_queue.qsize()}', file=sys.stderr) if batch is not None: out_queue.put(batch) else: @@ -486,7 +555,6 @@ def worker_manager(worker_result_queue, out_queue, num_workers, worker_manager_l break i += 1 - # print(f'[LocalWorkerManager] loaded {i} batches...', file=sys.stderr) worker_manager_lock.value = 0 else: if queue_size == prev_queue_size: @@ -494,7 +562,8 @@ def worker_manager(worker_result_queue, out_queue, num_workers, worker_manager_l prev_queue_size = queue_size time.sleep(0.2) - if finished: break + if finished: + break out_queue.put(None) @@ -507,7 +576,7 @@ def __init__(self, file_paths): assert isinstance(file_paths, str) self.file_paths = glob.glob(file_paths) - print(f'reading data files {self.file_paths}', file=sys.stderr) + print(f"reading data files {self.file_paths}", file=sys.stderr) example_num = 0 for _ in get_json_iterator_from_tar_file(self.file_paths): example_num += 1 @@ -519,8 +588,12 @@ def __len__(self): def __iter__(self): return self.get_iterator(progress=True) - def get_single_process_iterator(self, shuffle=False, progress=False) -> Iterable[Example]: - json_str_iter = get_json_iterator_from_tar_file(self.file_paths, shuffle, progress) + def get_single_process_iterator( + self, shuffle=False, progress=False + ) -> Iterable[Example]: + json_str_iter = get_json_iterator_from_tar_file( + self.file_paths, shuffle, progress + ) for json_str, meta in json_str_iter: tree_json_dict = json.loads(json_str) example = Example.from_json_dict(tree_json_dict, binary_file=meta) @@ -534,18 +607,30 @@ def _get_iterator(self, shuffle=False, num_workers=1): json_enc_queue = multiprocessing.Queue() example_queue = multiprocessing.Queue(maxsize=5000) - json_loader = multiprocessing.Process(target=json_line_reader, args=(self.file_paths, json_enc_queue, num_workers, - shuffle, False, None, False)) + json_loader = multiprocessing.Process( + target=json_line_reader, + args=( + self.file_paths, + json_enc_queue, + num_workers, + shuffle, + False, + None, + False, + ), + ) json_loader.daemon = True example_generators = [] for i in range(num_workers): - p = multiprocessing.Process(target=example_generator, - args=(json_enc_queue, example_queue, 1)) + p = multiprocessing.Process( + target=example_generator, args=(json_enc_queue, example_queue, 1) + ) p.daemon = True example_generators.append(p) json_loader.start() - for p in example_generators: p.start() + for p in example_generators: + p.start() num_finished_workers = 0 while True: @@ -554,44 +639,72 @@ def _get_iterator(self, shuffle=False, num_workers=1): yield example else: num_finished_workers += 1 - if num_finished_workers == num_workers: break + if num_finished_workers == num_workers: + break json_loader.join() - for p in example_generators: p.join() + for p in example_generators: + p.join() def get_iterator(self, shuffle=False, progress=True, num_workers=1): + iterator = self._get_iterator(shuffle, num_workers) if progress: - it_func = lambda x: tqdm(x, total=len(self), file=sys.stdout) - else: - it_func = lambda x: x - - return it_func(self._get_iterator(shuffle, num_workers)) - - def batch_iterator(self, batch_size: int, config: Dict, - return_examples=False, - return_prediction_target=None, - num_readers=3, - num_batchers=3, - progress=True, train=False, single_batcher=False) -> Iterable[Union[Batch, Dict[str, torch.Tensor]]]: + return tqdm(iterator, total=len(self), file=sys.stdout) + return iterator + + def batch_iterator( + self, + batch_size: int, + config: Dict, + return_examples=False, + return_prediction_target=None, + num_readers=3, + num_batchers=3, + progress=True, + train=False, + single_batcher=False, + ) -> Iterable[Union[Batch, Dict[str, torch.Tensor]]]: if progress: it_func = lambda x: tqdm(x, file=sys.stdout) else: it_func = lambda x: x if single_batcher: - return it_func(self._single_process_batch_iter(batch_size, config, num_readers, train)) + return it_func( + self._single_process_batch_iter(batch_size, config, num_readers, train) + ) else: - return it_func(self._batch_iterator(batch_size, config, num_readers, num_batchers, train, return_examples, return_prediction_target)) - - def _batch_iterator(self, batch_size: int, config: Dict, num_readers, num_batchers, train=False, return_examples=False, return_prediction_target=None) -> Iterable[Batch]: + return it_func( + self._batch_iterator( + batch_size, + config, + num_readers, + num_batchers, + train, + return_examples, + return_prediction_target, + ) + ) + + def _batch_iterator( + self, + batch_size: int, + config: Dict, + num_readers, + num_batchers, + train=False, + return_examples=False, + return_prediction_target=None, + ) -> Iterable[Batch]: global batcher_sync_msg - batcher_sync_msg = multiprocessing.Value('i', 0) + batcher_sync_msg = multiprocessing.Value("i", 0) json_enc_queue = multiprocessing.Queue(maxsize=10000) - worker_manager_lock = multiprocessing.Value('i', 0) + worker_manager_lock = multiprocessing.Value("i", 0) - json_loader = multiprocessing.Process(target=json_line_reader, - args=(self.file_paths, json_enc_queue, num_readers, - train, False)) + json_loader = multiprocessing.Process( + target=json_line_reader, + args=(self.file_paths, json_enc_queue, num_readers, train, False), + ) json_loader.daemon = True example_generators = [] worker_result_queue = torch_mp.Queue(maxsize=150) @@ -600,56 +713,77 @@ def _batch_iterator(self, batch_size: int, config: Dict, num_readers, num_batche return_prediction_target = train for i in range(num_readers): - p = multiprocessing.Process(target=example_to_batch, - args=(json_enc_queue, worker_result_queue, batch_size, train, config, worker_manager_lock, return_examples, return_prediction_target)) + p = multiprocessing.Process( + target=example_to_batch, + args=( + json_enc_queue, + worker_result_queue, + batch_size, + train, + config, + worker_manager_lock, + return_examples, + return_prediction_target, + ), + ) p.daemon = True example_generators.append(p) json_loader.start() - for p in example_generators: p.start() + for p in example_generators: + p.start() batch_queue = queue.Queue(maxsize=100) - worker_manager_thread = threading.Thread(target=worker_manager, args=(worker_result_queue, batch_queue, num_readers, worker_manager_lock, 100)) + worker_manager_thread = threading.Thread( + target=worker_manager, + args=( + worker_result_queue, + batch_queue, + num_readers, + worker_manager_lock, + 100, + ), + ) worker_manager_thread.start() while True: - # t1 = time.time() - main_process_queue_get_lock = 1 batch = batch_queue.get() - # print(f'[MainThread] local batch queue size {batch_queue.qsize()}', file=sys.stderr) - # print(f'{time.time() - t1} took to load a batch', file=sys.stderr) if batch is None: break else: yield batch worker_result_queue.close() - # print('start joining...') + # start joining... batcher_sync_msg.value = 1 json_loader.join() - # print('json_loader quitted') - for p in example_generators: p.join() + # json_loader quitted + for p in example_generators: + p.join() worker_manager_thread.join() - # print('example generators quitted') - # print('batchers quiteed') + # example generators quitted + # batchers quiteed sys.stdout.flush() sys.stderr.flush() - def _single_process_batch_iter(self, batch_size: int, config: Dict, num_readers=2, shuffle=False) -> Iterable[Batch]: + def _single_process_batch_iter( + self, batch_size: int, config: Dict, num_readers=2, shuffle=False + ) -> Iterable[Batch]: batcher = Batcher(config) - example_iter = self.get_iterator(shuffle=shuffle, progress=False, num_workers=num_readers) + example_iter = self.get_iterator( + shuffle=shuffle, progress=False, num_workers=num_readers + ) # t1 = time.time() batch_examples = [] batch_node_num = 0 # if example.ast.size < 300 and len(example.variable_name_map) > 0: - for example in filter(is_valid_example, example_iter): + for example in filter(is_valid_training_example, example_iter): batch_examples.append(example) batch_node_num += example.ast.size if batch_node_num >= batch_size: batch = batcher.to_batch(batch_examples) - # print(f'[Dataset] {time.time() - t1} took to load a batch', file=sys.stderr) yield batch batch_examples = [] @@ -661,7 +795,7 @@ def _single_process_batch_iter(self, batch_size: int, config: Dict, num_readers= yield batch -if __name__ == '__main__': - for _example in Dataset('data/0-trees.tar.gz'): +if __name__ == "__main__": + for _example in Dataset("data/0-trees.tar.gz"): if _example.ast.size > 200: print(_example.binary_file, _example.variable_name_map) diff --git a/neural-model/utils/eval_debin_prediction.py b/neural-model/utils/eval_debin_prediction.py index ad091a4..43842dd 100644 --- a/neural-model/utils/eval_debin_prediction.py +++ b/neural-model/utils/eval_debin_prediction.py @@ -1,6 +1,7 @@ # evaluate debin import json -from utils.dataset import get_json_iterator_from_tar_file, Example, Dataset + +from utils.dataset import Dataset, Example, get_json_iterator_from_tar_file from utils.evaluation import * @@ -10,7 +11,11 @@ def evaluate_debin(debin_output_path, test_example_path): debin_prediction_dict = dict() for example in debin_predicted_data.get_iterator(num_workers=5): - example_key = example.binary_file['file_name'].split('/')[-1] + '_' + str(example.ast.compilation_unit) + example_key = ( + example.binary_file["file_name"].split("/")[-1] + + "_" + + str(example.ast.compilation_unit) + ) assert example_key not in debin_prediction_dict @@ -38,12 +43,18 @@ def evaluate_debin(debin_output_path, test_example_path): test_set = Dataset(test_example_path) for example in test_set.get_iterator(num_workers=5): - example_key = example.binary_file['file_name'].split('/')[-1] + '_' + str(example.ast.compilation_unit) + example_key = ( + example.binary_file["file_name"].split("/")[-1] + + "_" + + str(example.ast.compilation_unit) + ) if example_key in debin_prediction_dict: rename_result = debin_prediction_dict[example_key] debin_predicted_func_num += 1 else: - rename_result = {k: k for k in example.variable_name_map} # use identity predictions + rename_result = { + k: k for k in example.variable_name_map + } # use identity predictions example_pred_accs = [] @@ -56,25 +67,29 @@ def evaluate_debin(debin_output_path, test_example_path): if gold_new_name != old_name: # and gold_new_name in model.vocab.target: need_rename_cases.append(var_metric) - if example.test_meta['function_name_in_train']: + if example.test_meta["function_name_in_train"]: func_name_in_train_acc_list.append(var_metric) else: func_name_not_in_train_acc_list.append(var_metric) - if example.test_meta['function_body_in_train']: + if example.test_meta["function_body_in_train"]: func_body_in_train_acc_list.append(var_metric) else: func_body_not_in_train_acc_list.append(var_metric) - eval_results = dict(corpus_need_rename_acc=Evaluator.average(need_rename_cases), - func_name_in_train_acc=Evaluator.average(func_name_in_train_acc_list), - func_name_not_in_train_acc=Evaluator.average(func_name_not_in_train_acc_list), - func_body_in_train_acc=Evaluator.average(func_body_in_train_acc_list), - func_body_not_in_train_acc=Evaluator.average(func_body_not_in_train_acc_list)) + eval_results = dict( + corpus_need_rename_acc=Evaluator.average(need_rename_cases), + func_name_in_train_acc=Evaluator.average(func_name_in_train_acc_list), + func_name_not_in_train_acc=Evaluator.average(func_name_not_in_train_acc_list), + func_body_in_train_acc=Evaluator.average(func_body_in_train_acc_list), + func_body_not_in_train_acc=Evaluator.average(func_body_not_in_train_acc_list), + ) - print('Num. functions debin predicted: ', debin_predicted_func_num) + print("Num. functions debin predicted: ", debin_predicted_func_num) print(eval_results) -if __name__ == '__main__': - evaluate_debin('data/debin.predictions.0.01.tar', 'data/all_trees_tokenized_0410/test.tar') +if __name__ == "__main__": + evaluate_debin( + "data/debin.predictions.0.01.tar", "data/all_trees_tokenized_0410/test.tar" + ) diff --git a/neural-model/utils/evaluation.py b/neural-model/utils/evaluation.py index 843af82..1d1f1fe 100644 --- a/neural-model/utils/evaluation.py +++ b/neural-model/utils/evaluation.py @@ -1,13 +1,11 @@ -from typing import Dict, List, Any +from typing import Any, Dict, List, Callable +import editdistance import numpy as np import torch - -import editdistance - +from model.model import RenamingModel from utils import nn_util from utils.dataset import Dataset -from model.model import RenamingModel class Evaluator(object): @@ -17,10 +15,9 @@ def get_soft_metrics(pred_name: str, gold_name: str) -> Dict: cer = float(edit_distance / len(gold_name)) acc = float(pred_name == gold_name) - return dict(edit_distance=edit_distance, - ref_len=len(gold_name), - cer=cer, - accuracy=acc) + return dict( + edit_distance=edit_distance, ref_len=len(gold_name), cer=cer, accuracy=acc + ) @staticmethod def average(metrics_list: List[Dict]) -> Dict: @@ -30,7 +27,9 @@ def average(metrics_list: List[Dict]) -> Dict: agg_results.setdefault(key, []).append(val) avg_results = dict() - avg_results['corpus_cer'] = sum(agg_results['edit_distance']) / sum(agg_results['ref_len']) + avg_results["corpus_cer"] = sum(agg_results["edit_distance"]) / sum( + agg_results["ref_len"] + ) for key, val in agg_results.items(): avg_results[key] = np.average(val) @@ -38,29 +37,39 @@ def average(metrics_list: List[Dict]) -> Dict: return avg_results @staticmethod - def evaluate_ppl(model: RenamingModel, dataset: Dataset, config: Dict, predicate: Any = None): - if predicate is None: - predicate = lambda e: True - - eval_batch_size = config['train']['batch_size'] - data_iter = dataset.batch_iterator(batch_size=eval_batch_size, - train=False, progress=True, - return_examples=False, - return_prediction_target=True, - config=model.config, - num_readers=config['train']['num_readers'], - num_batchers=config['train']['num_batchers']) + def _always_true(_: Any) -> bool: + return True + + @staticmethod + def evaluate_ppl( + model: RenamingModel, dataset: Dataset, config: Dict, predicate: Callable[[Any], bool] = _always_true + ): + eval_batch_size = config["train"]["batch_size"] + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=False, + return_prediction_target=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) was_training = model.training model.eval() - cum_log_probs = 0. + cum_log_probs = 0.0 cum_num_examples = 0 with torch.no_grad(): for batch in data_iter: nn_util.to(batch.tensor_dict, model.device) - result = model(batch.tensor_dict, batch.tensor_dict['prediction_target']) - log_probs = result['batch_log_prob'].cpu().tolist() - for e_id, test_meta in enumerate(batch.tensor_dict['test_meta']): + result = model( + batch.tensor_dict, batch.tensor_dict["prediction_target"] + ) + log_probs = result["batch_log_prob"].cpu().tolist() + for e_id, test_meta in enumerate(batch.tensor_dict["test_meta"]): if predicate(test_meta): log_prob = log_probs[e_id] cum_log_probs += log_prob @@ -74,15 +83,77 @@ def evaluate_ppl(model: RenamingModel, dataset: Dataset, config: Dict, predicate return ppl @staticmethod - def decode_and_evaluate(model: RenamingModel, dataset: Dataset, config: Dict, return_results=False, eval_batch_size=None): + def decode( + model: RenamingModel, dataset: Dataset, config: Dict, eval_batch_size=None + ): + if eval_batch_size is None: + if "eval_batch_size" in config["train"]: + eval_batch_size = config["train"]["eval_batch_size"] + else: + eval_batch_size = config["train"]["batch_size"] + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) + model.eval() + all_examples = dict() + + with torch.no_grad(): + for batch in data_iter: + examples = batch.examples + rename_results = model.predict(examples) + for example, rename_result in zip(examples, rename_results): + example_pred_accs = [] + top_rename_result = rename_result[0] + for old_name, gold_new_name in example.variable_name_map.items(): + pred = top_rename_result[old_name] + pred_new_name = pred["new_name"] + var_metric = Evaluator.get_soft_metrics( + pred_new_name, gold_new_name + ) + example_pred_accs.append(var_metric) + file_name = example.binary_file["file_name"] + line_num = example.binary_file["line_num"] + fun_name = example.ast.compilation_unit + all_examples[f"{file_name}_{line_num}_{fun_name}"] = ( + rename_result, + Evaluator.average(example_pred_accs), + ) + + return all_examples + + @staticmethod + def decode_and_evaluate( + model: RenamingModel, + dataset: Dataset, + config: Dict, + return_results=False, + eval_batch_size=None, + ): if eval_batch_size is None: - eval_batch_size = config['train']['eval_batch_size'] if 'eval_batch_size' in config['train'] else config['train']['batch_size'] - data_iter = dataset.batch_iterator(batch_size=eval_batch_size, - train=False, progress=True, - return_examples=True, - config=model.config, - num_readers=config['train']['num_readers'], - num_batchers=config['train']['num_batchers']) + if "eval_batch_size" in config["train"]: + eval_batch_size = config["train"]["eval_batch_size"] + else: + eval_batch_size = config["train"]["batch_size"] + + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) was_training = model.training model.eval() @@ -90,10 +161,10 @@ def decode_and_evaluate(model: RenamingModel, dataset: Dataset, config: Dict, re variable_acc_list = [] need_rename_cases = [] - func_name_in_train_acc_list = [] - func_name_not_in_train_acc_list = [] - func_body_in_train_acc_list = [] - func_body_not_in_train_acc_list = [] + func_name_in_train_acc = [] + func_name_not_in_train_acc = [] + func_body_in_train_acc = [] + func_body_not_in_train_acc = [] all_examples = dict() @@ -107,30 +178,39 @@ def decode_and_evaluate(model: RenamingModel, dataset: Dataset, config: Dict, re top_rename_result = rename_result[0] for old_name, gold_new_name in example.variable_name_map.items(): pred = top_rename_result[old_name] - pred_new_name = pred['new_name'] - var_metric = Evaluator.get_soft_metrics(pred_new_name, gold_new_name) + pred_new_name = pred["new_name"] + var_metric = Evaluator.get_soft_metrics( + pred_new_name, gold_new_name + ) # is_correct = pred_new_name == gold_new_name example_pred_accs.append(var_metric) - if gold_new_name != old_name: # and gold_new_name in model.vocab.target: + if gold_new_name != old_name: need_rename_cases.append(var_metric) - if example.test_meta['function_name_in_train']: - func_name_in_train_acc_list.append(var_metric) + if example.test_meta["function_name_in_train"]: + func_name_in_train_acc.append(var_metric) else: - func_name_not_in_train_acc_list.append(var_metric) + func_name_not_in_train_acc.append(var_metric) - if example.test_meta['function_body_in_train']: - func_body_in_train_acc_list.append(var_metric) + if example.test_meta["function_body_in_train"]: + func_body_in_train_acc.append(var_metric) else: - func_body_not_in_train_acc_list.append(var_metric) + func_body_not_in_train_acc.append(var_metric) variable_acc_list.extend(example_pred_accs) example_acc_list.append(example_pred_accs) if return_results: - all_examples[example.binary_file['file_name'] + '_' + str(example.binary_file['line_num'])] = (rename_result, Evaluator.average(example_pred_accs)) - # all_examples.append((example, rename_result, example_pred_accs)) + example_key = ( + f"{example.binary_file['file_name']}_" + f"{example.binary_file['line_num']}" + ) + + all_examples[example_key] = ( + rename_result, + Evaluator.average(example_pred_accs), + ) valid_example_num = len(example_acc_list) num_variables = len(variable_acc_list) @@ -139,19 +219,26 @@ def decode_and_evaluate(model: RenamingModel, dataset: Dataset, config: Dict, re if was_training: model.train() - eval_results = dict(corpus_acc=corpus_acc, - corpus_need_rename_acc=Evaluator.average(need_rename_cases), - func_name_in_train_acc=Evaluator.average(func_name_in_train_acc_list), - func_name_not_in_train_acc=Evaluator.average(func_name_not_in_train_acc_list), - func_body_in_train_acc=Evaluator.average(func_body_in_train_acc_list), - func_body_not_in_train_acc=Evaluator.average(func_body_not_in_train_acc_list), - num_variables=num_variables, - num_valid_examples=valid_example_num) + need_rename_acc = Evaluator.average(need_rename_cases) + name_in_train_acc = Evaluator.average(func_name_in_train_acc) + name_not_in_train_acc = Evaluator.average(func_name_not_in_train_acc) + body_in_train_acc = Evaluator.average(func_body_in_train_acc) + body_not_in_train_acc = Evaluator.average(func_body_not_in_train_acc) + eval_results = dict( + corpus_acc=corpus_acc, + corpus_need_rename_acc=need_rename_acc, + func_name_in_train_acc=name_in_train_acc, + func_name_not_in_train_acc=name_not_in_train_acc, + func_body_in_train_acc=body_in_train_acc, + func_body_not_in_train_acc=body_not_in_train_acc, + num_variables=num_variables, + num_valid_examples=valid_example_num, + ) if return_results: return eval_results, all_examples return eval_results -if __name__ == '__main__': - print(Evaluator.get_soft_metrics('file_name', 'filename')) +if __name__ == "__main__": + print(Evaluator.get_soft_metrics("file_name", "filename")) diff --git a/neural-model/utils/get_stat.py b/neural-model/utils/get_stat.py index d1d7cb4..f403279 100644 --- a/neural-model/utils/get_stat.py +++ b/neural-model/utils/get_stat.py @@ -1,15 +1,15 @@ # compute training data statistics -from collections import Counter import gc -import numpy as np +from collections import Counter +import numpy as np from utils.dataset import Dataset gc.collect() def compute_dataset_stat(): - train_path = 'data/all_trees_tokenized_0410/train-shard-*.tar' + train_path = "data/all_trees_tokenized_0410/train-shard-*.tar" train_set = Dataset(train_path) num_train_examples = 0 @@ -18,8 +18,8 @@ def compute_dataset_stat(): ast_sizes = [] var_name_freq = Counter() - dev_set = Dataset('data/all_trees_tokenized_0410/dev.tar') - test_set = Dataset('data/all_trees_tokenized_0410/test.tar') + dev_set = Dataset("data/all_trees_tokenized_0410/dev.tar") + test_set = Dataset("data/all_trees_tokenized_0410/test.tar") for dataset_id, dataset in enumerate([train_set, dev_set, test_set]): for example in dataset.get_iterator(num_workers=5): @@ -39,17 +39,21 @@ def compute_dataset_stat(): del example - with open('data/all_trees_tokenized_0410/train_var_name_freq.txt', 'w') as f: + with open("data/all_trees_tokenized_0410/train_var_name_freq.txt", "w") as f: for var_name, freq in var_name_freq.most_common(): - f.write(f'{var_name}\t{freq}\n') + f.write(f"{var_name}\t{freq}\n") - with open('data/all_trees_tokenized_0410/stat.txt', 'w') as f: + with open("data/all_trees_tokenized_0410/stat.txt", "w") as f: print(len(train_set), len(dev_set), len(test_set), file=f) - print('num total functions: ', num_train_examples, file=f) - print('avg. size of AST:', np.average(ast_sizes), file=f) - print('avg. num of variables:', np.average(num_variables), file=f) - print('avg. num of variables with new name:', np.average(num_variables_with_new_name), file=f) + print("num total functions: ", num_train_examples, file=f) + print("avg. size of AST:", np.average(ast_sizes), file=f) + print("avg. num of variables:", np.average(num_variables), file=f) + print( + "avg. num of variables with new name:", + np.average(num_variables_with_new_name), + file=f, + ) -if __name__ == '__main__': +if __name__ == "__main__": compute_dataset_stat() diff --git a/neural-model/utils/grammar.py b/neural-model/utils/grammar.py index 3f0a017..54ec9f3 100644 --- a/neural-model/utils/grammar.py +++ b/neural-model/utils/grammar.py @@ -1,16 +1,18 @@ class Grammar(object): def __init__(self, syntax_types, variable_types): self.syntax_types = list(sorted(syntax_types)) - self.variable_types = ['', ''] + list(sorted(variable_types)) + self.variable_types = ["", ""] + list(sorted(variable_types)) self.syntax_type_to_id = {type: id for id, type in enumerate(self.syntax_types)} - self._variable_type_to_id = {type: id for id, type in enumerate(self.variable_types)} + self._variable_type_to_id = { + type: id for id, type in enumerate(self.variable_types) + } def variable_type_to_id(self, type_token): if type_token in self._variable_type_to_id: return self._variable_type_to_id[type_token] - return self._variable_type_to_id[''] + return self._variable_type_to_id[""] @property def params(self): @@ -18,6 +20,6 @@ def params(self): @classmethod def load(cls, params): - grammar = cls(params['syntax_types'], params['variable_types']) + grammar = cls(params["syntax_types"], params["variable_types"]) return grammar diff --git a/neural-model/utils/graph.py b/neural-model/utils/graph.py index 7b25835..23f3cae 100644 --- a/neural-model/utils/graph.py +++ b/neural-model/utils/graph.py @@ -1,5 +1,5 @@ -from collections import defaultdict, OrderedDict -from typing import List, Dict +from collections import OrderedDict, defaultdict +from typing import Dict, List from utils.ast import AbstractSyntaxTree, SyntaxNode @@ -24,7 +24,9 @@ def register_tree(self, ast: AbstractSyntaxTree): for node in ast: self.register_node(ast_id, node.node_id) - def register_node(self, tree_id, node, group='ast_nodes', return_node_index_in_group=False): + def register_node( + self, tree_id, node, group="ast_nodes", return_node_index_in_group=False + ): if group not in self.node_groups[tree_id]: self.node_groups[tree_id][group] = OrderedDict() @@ -40,7 +42,7 @@ def register_node(self, tree_id, node, group='ast_nodes', return_node_index_in_g return packed_node_id - def get_packed_node_id(self, tree_id, node, group='ast_nodes'): + def get_packed_node_id(self, tree_id, node, group="ast_nodes"): if isinstance(node, SyntaxNode): node = node.node_id return self.node_groups[tree_id][group][node] @@ -55,9 +57,10 @@ def nodes(self): def get_nodes_by_group(self, group): for i in range(self.tree_num): - for node, packed_node_id in self.node_groups[i][group].items(): + node_group = self.node_groups[i] + for node, packed_node_id in node_group[group].items(): yield node, packed_node_id @property def tree_num(self): - return len(self.trees) \ No newline at end of file + return len(self.trees) diff --git a/neural-model/utils/gz_to_jsonl.py b/neural-model/utils/gz_to_jsonl.py index 22c131c..0858e97 100644 --- a/neural-model/utils/gz_to_jsonl.py +++ b/neural-model/utils/gz_to_jsonl.py @@ -1,16 +1,17 @@ -import tarfile -import sys import json +import sys +import tarfile + from tqdm import tqdm file_name = sys.argv[1] tgt_file_name = sys.argv[2] -with open(tgt_file_name, 'w') as f_tgt: - with tarfile.open(file_name, 'r') as f: +with open(tgt_file_name, "w") as f_tgt: + with tarfile.open(file_name, "r") as f: for filename in tqdm(f.getmembers()): json_file = f.extractfile(filename) if json_file is not None: json_str = json_file.read() json_str = json.dumps(json.loads(json_str)) - f_tgt.write(json_str + '\n') + f_tgt.write(json_str + "\n") diff --git a/neural-model/utils/lexer.py b/neural-model/utils/lexer.py index bdf7acd..4ca4d4b 100644 --- a/neural-model/utils/lexer.py +++ b/neural-model/utils/lexer.py @@ -7,10 +7,10 @@ # reference shadowed global variables. We throw away this operator. from enum import Enum, auto + from pygments import lex -from pygments.token import Token -from pygments.token import is_token_subtype from pygments.lexers.c_cpp import CLexer, inherit +from pygments.token import Token, is_token_subtype Token.Placeholder = Token.Token.Placeholder @@ -56,14 +56,14 @@ def get_tokens(self, var_names=Names.RAW): elif is_token_subtype(token_type, Token.Comment): continue # Skip the :: token added by HexRays - elif is_token_subtype(token_type, Token.Operator) and token == '::': + elif is_token_subtype(token_type, Token.Operator) and token == "::": continue # Replace the text of placeholder tokens elif is_token_subtype(token_type, Token.Placeholder): yield { - Names.RAW : (token_type, token), - Names.SOURCE : (token_type, token.split('@@')[2]), - Names.TARGET : (token_type, token.split('@@')[3]), + Names.RAW: (token_type, token), + Names.SOURCE: (token_type, token.split("@@")[2]), + Names.TARGET: (token_type, token.split("@@")[3]), }[var_names] elif not is_token_subtype(token_type, Token.Text): yield (token_type, token.strip()) @@ -77,37 +77,45 @@ def get_tokens(self, var_names=Names.RAW): class HexRaysLexer(CLexer): # Additional tokens tokens = { - 'statements' : [ - (r'->', Token.Operator), - (r'\+\+', Token.Operator), - (r'--', Token.Operator), - (r'==', Token.Operator), - (r'!=', Token.Operator), - (r'>=', Token.Operator), - (r'<=', Token.Operator), - (r'&&', Token.Operator), - (r'\|\|', Token.Operator), - (r'\+=', Token.Operator), - (r'-=', Token.Operator), - (r'\*=', Token.Operator), - (r'/=', Token.Operator), - (r'%=', Token.Operator), - (r'&=', Token.Operator), - (r'\^=', Token.Operator), - (r'\|=', Token.Operator), - (r'<<=', Token.Operator), - (r'>>=', Token.Operator), - (r'<<', Token.Operator), - (r'>>', Token.Operator), - (r'\.\.\.', Token.Operator), - (r'##', Token.Operator), - (r'::', Token.Operator), - (r'@@VAR_[0-9]+@@\w+@@\w+', Token.Placeholder.Var), - inherit + "statements": [ + (r"->", Token.Operator), + (r"\+\+", Token.Operator), + (r"--", Token.Operator), + (r"==", Token.Operator), + (r"!=", Token.Operator), + (r">=", Token.Operator), + (r"<=", Token.Operator), + (r"&&", Token.Operator), + (r"\|\|", Token.Operator), + (r"\+=", Token.Operator), + (r"-=", Token.Operator), + (r"\*=", Token.Operator), + (r"/=", Token.Operator), + (r"%=", Token.Operator), + (r"&=", Token.Operator), + (r"\^=", Token.Operator), + (r"\|=", Token.Operator), + (r"<<=", Token.Operator), + (r">>=", Token.Operator), + (r"<<", Token.Operator), + (r">>", Token.Operator), + (r"\.\.\.", Token.Operator), + (r"##", Token.Operator), + (r"::", Token.Operator), + (r"@@VAR_[0-9]+@@\w+@@\w+", Token.Placeholder.Var), + inherit, ] } -if __name__ == '__main__': - code = '__int64 (__fastcall **)(unsigned __int64, signed __int64, __int64, _QWORD, __int64, signed __int64) {int a = "asdfsdf"; int b =123; a = asd::safd() sadf=12 /*asdf*/}' - print([str(token[0]) for token in Lexer(code).get_tokens()]) \ No newline at end of file +if __name__ == "__main__": + code = """ +__int64 (__fastcall **)(unsigned __int64, signed __int64, + __int64, _QWORD, __int64, signed __int64) { + int a = "asdfsdf"; + int b =123; + a = asd::safd(); + sadf=12 + /*asdf*/ +}""" + print([str(token[0]) for token in Lexer(code).get_tokens()]) diff --git a/neural-model/utils/nn_util.py b/neural-model/utils/nn_util.py index 99ce990..630a729 100644 --- a/neural-model/utils/nn_util.py +++ b/neural-model/utils/nn_util.py @@ -4,7 +4,6 @@ import numpy as np import torch - SMALL_NUMBER = 1e-8 @@ -15,8 +14,8 @@ def glorot_init(params): def to(data, device: torch.device): - if 'adj_lists' in data: - [x.to(device) for x in data['adj_lists']] + if "adj_lists" in data: + [x.to(device) for x in data["adj_lists"]] if isinstance(data, dict): for key, val in data.items(): @@ -44,7 +43,7 @@ def batch_iter(data, batch_size, shuffle=False): np.random.shuffle(index_array) for i in range(batch_num): - indices = index_array[i * batch_size: (i + 1) * batch_size] + indices = index_array[i * batch_size : (i + 1) * batch_size] examples = [data[idx] for idx in indices] yield examples @@ -65,23 +64,25 @@ def get_tensor_dict_size(tensor_dict): return total_num_elements -def dot_prod_attention(h_t: torch.Tensor, - src_encoding: torch.Tensor, - src_encoding_att_linear: torch.Tensor, - mask: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: - # (batch_size, src_sent_len) - att_weight = torch.bmm(src_encoding_att_linear, h_t.unsqueeze(2)).squeeze(2) +def dot_prod_attention( + h_t: torch.Tensor, + src_encoding: torch.Tensor, + src_encoding_att_linear: torch.Tensor, + mask: torch.Tensor = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + # (batch_size, src_sent_len) + att_weight = torch.bmm(src_encoding_att_linear, h_t.unsqueeze(2)).squeeze(2) - if mask is not None: - att_weight.data.masked_fill_((1. - mask).bool(), -float('inf')) + if mask is not None: + att_weight.data.masked_fill_((1.0 - mask).bool(), -float("inf")) - softmaxed_att_weight = torch.softmax(att_weight, dim=-1) + softmaxed_att_weight = torch.softmax(att_weight, dim=-1) - att_view = (att_weight.size(0), 1, att_weight.size(1)) - # (batch_size, hidden_size) - ctx_vec = torch.bmm(softmaxed_att_weight.view(*att_view), src_encoding).squeeze(1) + att_view = (att_weight.size(0), 1, att_weight.size(1)) + # (batch_size, hidden_size) + ctx_vec = torch.bmm(softmaxed_att_weight.view(*att_view), src_encoding).squeeze(1) - return ctx_vec, softmaxed_att_weight + return ctx_vec, softmaxed_att_weight def get_lengths_from_binary_sequence_mask(mask: torch.Tensor): @@ -114,21 +115,27 @@ def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor): Returns ------- sorted_tensor : torch.FloatTensor - The original tensor sorted along the batch dimension with respect to sequence_lengths. + The original tensor sorted along the batch dimension with respect to + sequence_lengths. sorted_sequence_lengths : torch.LongTensor The original sequence_lengths sorted by decreasing size. restoration_indices : torch.LongTensor Indices into the sorted_tensor such that ``sorted_tensor.index_select(0, restoration_indices) == original_tensor`` permuation_index : torch.LongTensor - The indices used to sort the tensor. This is useful if you want to sort many - tensors using the same ordering. + The indices used to sort the tensor. This is useful if you want to sort + many tensors using the same ordering. + """ - if not isinstance(tensor, torch.Tensor) or not isinstance(sequence_lengths, torch.Tensor): + if not isinstance(tensor, torch.Tensor) or not isinstance( + sequence_lengths, torch.Tensor + ): raise ValueError("Both the tensor and sequence lengths must be torch.Tensors.") - sorted_sequence_lengths, permutation_index = sequence_lengths.sort(0, descending=True) + sorted_sequence_lengths, permutation_index = sequence_lengths.sort( + 0, descending=True + ) sorted_tensor = tensor.index_select(0, permutation_index) index_range = torch.arange(0, len(sequence_lengths), device=sequence_lengths.device) @@ -136,4 +143,9 @@ def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor): # sequence lengths and returning the now sorted indices. _, reverse_mapping = permutation_index.sort(0, descending=False) restoration_indices = index_range.index_select(0, reverse_mapping) - return sorted_tensor, sorted_sequence_lengths, restoration_indices, permutation_index + return ( + sorted_tensor, + sorted_sequence_lengths, + restoration_indices, + permutation_index, + ) diff --git a/neural-model/utils/preprocess.py b/neural-model/utils/preprocess.py index 3d1a707..1fc3ce3 100644 --- a/neural-model/utils/preprocess.py +++ b/neural-model/utils/preprocess.py @@ -12,64 +12,106 @@ import glob import multiprocessing +import os import random +import sys import tarfile from collections import Iterable +from multiprocessing import Process from typing import Tuple -import ujson as json -from docopt import docopt -import os, sys -from multiprocessing import Process import numpy as np - +import ujson as json +from docopt import docopt +from tqdm import tqdm from utils.ast import SyntaxNode -from utils.code_processing import canonicalize_code, annotate_type, canonicalize_constants, preprocess_ast, \ - tokenize_raw_code +from utils.code_processing import ( + annotate_type, + canonicalize_code, + canonicalize_constants, + preprocess_ast, + tokenize_raw_code, +) from utils.dataset import Example, json_line_reader -from tqdm import tqdm all_functions = dict() # indexed by binaries def is_valid_example(example): try: - is_valid = example.ast.size < 300 and \ - len(example.variable_name_map) > 0 and \ - any(k != v for k, v in example.variable_name_map.items()) + is_valid = ( + example.ast.size < 300 + and len(example.variable_name_map) > 0 + and any(k != v for k, v in example.variable_name_map.items()) + ) except RecursionError: is_valid = False return is_valid +def generate_example(json_str, binary_file): + tree_json_dict = json.loads(json_str) + + root = SyntaxNode.from_json_dict(tree_json_dict["ast"]) + + preprocess_ast(root, code=tree_json_dict["raw_code"]) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) + tree_json_dict["code_tokens"] = code_tokens + + # add function name to the name field of the root block + root.name = tree_json_dict["function"] + root.named_fields.add("name") + + new_json_dict = root.to_json_dict() + tree_json_dict["ast"] = new_json_dict + json_str = json.dumps(tree_json_dict) + + example = Example.from_json_dict( + tree_json_dict, + binary_file=binary_file, + json_str=json_str, + code_tokens=code_tokens, + ) + + if True or is_valid_example(example): + canonical_code = canonicalize_code(example.ast.code) + example.canonical_code = canonical_code + return example + else: + return None + + def example_generator(json_queue, example_queue, args, consumer_num=1): - enable_filter = not args['--no-filtering'] + enable_filter = not args["--no-filtering"] while True: payload = json_queue.get() - if payload is None: break + if payload is None: + break examples = [] for json_str, meta in payload: tree_json_dict = json.loads(json_str) - root = SyntaxNode.from_json_dict(tree_json_dict['ast']) + root = SyntaxNode.from_json_dict(tree_json_dict["ast"]) # root_reconstr = SyntaxNode.from_json_dict(root.to_json_dict()) # assert root == root_reconstr - preprocess_ast(root, code=tree_json_dict['raw_code']) - code_tokens = tokenize_raw_code(tree_json_dict['raw_code']) - tree_json_dict['code_tokens'] = code_tokens + preprocess_ast(root, code=tree_json_dict["raw_code"]) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) + tree_json_dict["code_tokens"] = code_tokens # add function name to the name field of the root block - root.name = tree_json_dict['function'] - root.named_fields.add('name') + root.name = tree_json_dict["function"] + root.named_fields.add("name") new_json_dict = root.to_json_dict() - tree_json_dict['ast'] = new_json_dict + tree_json_dict["ast"] = new_json_dict json_str = json.dumps(tree_json_dict) - example = Example.from_json_dict(tree_json_dict, binary_file=meta, json_str=json_str) + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, json_str=json_str + ) is_valid = is_valid_example(example) if enable_filter else True if is_valid: @@ -82,39 +124,51 @@ def example_generator(json_queue, example_queue, args, consumer_num=1): for i in range(consumer_num): example_queue.put(None) - print('example generator quited!') + print("example generator quit!") def main(args): - np.random.seed(1234) - random.seed(1992) + test_run = bool(args["--test-file"]) + if test_run: + np.random.seed(1234) + random.seed(1992) - tgt_folder = args['TARGET_FOLDER'] - pattern_list = args['TAR_FILES'].split(',') + tgt_folder = args["TARGET_FOLDER"] + pattern_list = args["TAR_FILES"].split(",") tar_files = [] for pattern in pattern_list: tar_files.extend(glob.glob(pattern)) print(tar_files) - shard_size = int(args['--shard-size']) os.system(f'mkdir -p "{tgt_folder}/files"') num_workers = 14 for tar_file in tar_files: - print(f'read {tar_file}') + print(f"read {tar_file}") valid_example_count = 0 json_enc_queue = multiprocessing.Queue() example_queue = multiprocessing.Queue(maxsize=2000) - json_loader = multiprocessing.Process(target=json_line_reader, - args=(os.path.expanduser(tar_file), json_enc_queue, num_workers, False, False, 'binary_file')) + json_loader = multiprocessing.Process( + target=json_line_reader, + args=( + os.path.expanduser(tar_file), + json_enc_queue, + num_workers, + False, + False, + "binary_file", + ), + ) json_loader.daemon = True json_loader.start() example_generators = [] for i in range(num_workers): - p = multiprocessing.Process(target=example_generator, args=(json_enc_queue, example_queue, args, 1)) + p = multiprocessing.Process( + target=example_generator, args=(json_enc_queue, example_queue, args, 1) + ) p.daemon = True p.start() example_generators.append(p) @@ -123,85 +177,46 @@ def main(args): while True: payload = example_queue.get() if payload is None: - print('received None!') + print("received None!") n_finished += 1 - if n_finished == num_workers: break + if n_finished == num_workers: + break continue examples = payload if examples: - json_file_name = examples[0].binary_file['file_name'].split('/')[-1] - with open(os.path.join(tgt_folder, 'files/', json_file_name), 'w') as f: + json_file_name = examples[0].binary_file["file_name"].split("/")[-1] + with open(os.path.join(tgt_folder, "files/", json_file_name), "w") as f: for example in examples: - f.write(example.json_str + '\n') - all_functions.setdefault(json_file_name, dict())[example.ast.compilation_unit] = example.canonical_code + f.write(example.json_str + "\n") + all_functions.setdefault(json_file_name, dict())[ + example.ast.compilation_unit + ] = example.canonical_code valid_example_count += len(examples) - print('valid examples: ', valid_example_count) + print("valid examples: ", valid_example_count) json_enc_queue.close() example_queue.close() json_loader.join() - for p in example_generators: p.join() + for p in example_generators: + p.join() cur_dir = os.getcwd() - all_files = glob.glob(os.path.join(tgt_folder, 'files/*.jsonl')) - file_prefix = os.path.join(tgt_folder, 'files/') + all_files = glob.glob(os.path.join(tgt_folder, "files/*.jsonl")) sorted(all_files) # sort all files by names all_files = list(all_files) file_num = len(all_files) - print('Total valid binary file num: ', file_num) - - test_file = args['--test-file'] - if test_file: - print(f'using test file {test_file}') - with tarfile.open(test_file, 'r') as f: - test_files = [os.path.join(file_prefix, x.name.split('/')[-1]) for x in f.getmembers() if x.name.endswith('.jsonl')] - dev_file_num = 0 - else: - print(f'randomly sample test file {test_file}') - test_file_num = int(file_num * 0.1) - dev_file_num = int(file_num * 0.1) - test_files = list(np.random.choice(all_files, size=test_file_num, replace=False)) - - test_files_set = set(test_files) - train_files = [fname for fname in all_files if fname not in test_files_set] - - if dev_file_num == 0: - dev_file_num = int(len(train_files) * 0.1) - - np.random.shuffle(train_files) - dev_files = train_files[-dev_file_num:] - train_files = train_files[:-dev_file_num] - - train_functions = dict() - for train_file in train_files: - file_name = train_file.split('/')[-1] - for func_name, func in all_functions[file_name].items(): - train_functions.setdefault(func_name, set()).add(func) - - print(f'number training: {len(train_files)}, number dev: {len(dev_files)}, number test: {len(test_files)}') - print('dump training files') - shards = [train_files[i:i + shard_size] for i in range(0, len(train_files), shard_size)] - for shard_id, shard_files in enumerate(shards): - print(f'Preparing shard {shard_id}, {len(shard_files)} files: ') - with open(os.path.join(tgt_folder, 'file_list.txt'), 'w') as f: - for file_name in shard_files: - f.write(file_name.split('/')[-1] + '\n') - - os.chdir(os.path.join(tgt_folder, 'files')) - print('creating tar file...') - os.system(f'tar cf ../train-shard-{shard_id}.tar -T ../file_list.txt') - os.chdir(cur_dir) + print(f"{file_num} valid binary files.") - def _dump_dev_file(tgt_file_name, file_names): - with open(os.path.join(tgt_folder, 'file_list.txt'), 'w') as f: + def _dump_file(tgt_file_name, file_names, dev=False): + with open(os.path.join(tgt_folder, "file_list.txt"), "w") as f: for file_name in file_names: - last_file_name = file_name.split('/')[-1] - f.write(last_file_name + '\n') + last_file_name = file_name.split("/")[-1] + f.write(last_file_name + "\n") with open(file_name) as fr: all_lines = fr.readlines() @@ -209,35 +224,44 @@ def _dump_dev_file(tgt_file_name, file_names): replace_lines = [] for line in all_lines: json_dict = json.loads(line.strip()) - func_name = json_dict['function'] - canonical_code = all_functions[last_file_name][func_name] - func_name_in_train = False - func_body_in_train = False - if func_name in train_functions: - func_name_in_train = True - if canonical_code in train_functions[func_name]: - func_body_in_train = True - - json_dict['test_meta'] = dict(function_name_in_train=func_name_in_train, - function_body_in_train=func_body_in_train) + if dev: + func_name = json_dict["function"] + canonical_code = all_functions[last_file_name][func_name] + func_name_in_train = False + func_body_in_train = False + if func_name in train_functions: + func_name_in_train = True + if canonical_code in train_functions[func_name]: + func_body_in_train = True + + json_dict["test_meta"] = dict( + function_name_in_train=func_name_in_train, + function_body_in_train=func_body_in_train, + ) new_json_str = json.dumps(json_dict) replace_lines.append(new_json_str.strip()) - with open(file_name, 'w') as fw: + with open(file_name, "w") as fw: for line in replace_lines: - fw.write(line + '\n') + fw.write(line + "\n") - os.chdir(os.path.join(tgt_folder, 'files')) - print('creating tar file...') - os.system(f'tar cf ../{tgt_file_name} -T ../file_list.txt') + os.chdir(os.path.join(tgt_folder, "files")) + print("creating tar file...") + os.system(f"tar cf ../{tgt_file_name} -T ../file_list.txt") os.chdir(cur_dir) - print('dump dev files') - _dump_dev_file('dev.tar', dev_files) - print('dump test files') - _dump_dev_file('test.tar', test_files) + if test_run: + print("dumping dev files") + _dump_dev_file("dev.tar", dev_files, dev=True) + + print("dumping test files") + _dump_dev_file("test.tar", test_files, dev=True) + else: + print("dumping files") + _dump_file("preprocessed.tar", all_files) + -if __name__ == '__main__': +if __name__ == "__main__": args = docopt(__doc__) main(args) diff --git a/neural-model/utils/sequential_preprocess.py b/neural-model/utils/sequential_preprocess.py index 35e526b..c19748a 100644 --- a/neural-model/utils/sequential_preprocess.py +++ b/neural-model/utils/sequential_preprocess.py @@ -7,37 +7,39 @@ -h --help Show this screen. """ +import ujson as json from docopt import docopt from utils.code_processing import VARIABLE_ANNOTATION -from utils.lexer import * from utils.dataset import Dataset -import ujson as json +from utils.lexer import * def processor(input_queue, output_queue): while True: payload = input_queue.get() - if payload is None: break + if payload is None: + break examples = [] for json_str, meta in payload: tree_json_dict = json.loads(json_str) - code_tokens = tokenize_raw_code(tree_json_dict['raw_code']) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) - tree_json_dict['code_tokens'] = code_tokens + tree_json_dict["code_tokens"] = code_tokens json_str = json.dumps(tree_json_dict) - preprocess_ast(root, code=tree_json_dict['raw_code']) + preprocess_ast(root, code=tree_json_dict["raw_code"]) # add function name to the name field of the root block - root.name = tree_json_dict['function'] - root.named_fields.add('name') + root.name = tree_json_dict["function"] + root.named_fields.add("name") new_json_dict = root.to_json_dict() - tree_json_dict['ast'] = new_json_dict - + tree_json_dict["ast"] = new_json_dict - example = Example.from_json_dict(tree_json_dict, binary_file=meta, json_str=json_str) + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, json_str=json_str + ) if is_valid_example(example): canonical_code = canonicalize_code(example.ast.code) @@ -49,12 +51,12 @@ def processor(input_queue, output_queue): for i in range(consumer_num): example_queue.put(None) - print('example generator quited!') + print("example generator quited!") def main(args): - dataset = Dataset(args['TAR_FILES']) - code_line_file = open(args['OUTPUT_CODE_FILE'], 'w') + dataset = Dataset(args["TAR_FILES"]) + code_line_file = open(args["OUTPUT_CODE_FILE"], "w") all_preserved_tokens = set() for example in dataset.get_iterator(num_workers=5): code = example.ast.code @@ -66,11 +68,11 @@ def main(args): code_line_file.close() - with open(args['OUTPUT_CODE_FILE'] + '.preserved_tokens.txt', 'w') as f: + with open(args["OUTPUT_CODE_FILE"] + ".preserved_tokens.txt", "w") as f: for token in all_preserved_tokens: - f.write(token + '\n') + f.write(token + "\n") -if __name__ == '__main__': +if __name__ == "__main__": args = docopt(__doc__) main(args) diff --git a/neural-model/utils/subsample.py b/neural-model/utils/subsample.py index bc75533..df672da 100644 --- a/neural-model/utils/subsample.py +++ b/neural-model/utils/subsample.py @@ -1,7 +1,9 @@ -import os, sys import glob -from sh import tar +import os +import sys + import numpy as np +from sh import tar def main(): @@ -14,29 +16,36 @@ def main(): print(tar_files) work_dir = os.path.dirname(tar_files[0]) - os.system(f'mkdir -p {work_dir}/binary_files/') + os.system(f"mkdir -p {work_dir}/binary_files/") for tar_file in tar_files: - print(f'extracting {tar_file}') - tar('-C', f'{work_dir}/binary_files/', '-xvf', tar_file) + print(f"extracting {tar_file}") + tar("-C", f"{work_dir}/binary_files/", "-xvf", tar_file) - all_files = glob.glob(f'{work_dir}/binary_files/*.jsonl') + all_files = glob.glob(f"{work_dir}/binary_files/*.jsonl") all_files.sort() - print(f'{len(all_files)} in total') - sampled_files = np.random.choice(all_files, replace=False, size=int(sample_ratio * len(all_files))) - print(f'{len(sampled_files)} sampled files') + print(f"{len(all_files)} in total") + sampled_files = np.random.choice( + all_files, replace=False, size=int(sample_ratio * len(all_files)) + ) + print(f"{len(sampled_files)} sampled files") os.chdir(work_dir) - with open(f'sampled_binaries.txt', 'w') as f: + with open(f"sampled_binaries.txt", "w") as f: for fname in sampled_files: fname = os.path.basename(fname) - f.write(fname + '\n') + f.write(fname + "\n") - print('creating tar file') - os.chdir('binary_files/') - tar('-cf', f'../sampled_binaries_{sample_ratio}.tar', '-T', '../sampled_binaries.txt') + print("creating tar file") + os.chdir("binary_files/") + tar( + "-cf", + f"../sampled_binaries_{sample_ratio}.tar", + "-T", + "../sampled_binaries.txt", + ) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/neural-model/utils/util.py b/neural-model/utils/util.py index d5d08e1..8b01106 100644 --- a/neural-model/utils/util.py +++ b/neural-model/utils/util.py @@ -1,5 +1,9 @@ -import psutil, gc, sys, os import collections +import gc +import os +import sys + +import psutil import torch @@ -9,11 +13,11 @@ class cached_property(object): itself with an ordinary attribute. Deleting the attribute resets the property. - Source: https://github.com/bottlepy/bottle/commit/fa7733e075da0d790d809aa3d2f53071897e6f76 + Source: https://github.com/bottlepy/bottle/commit/fa7733 """ def __init__(self, func): - self.__doc__ = getattr(func, '__doc__') + self.__doc__ = getattr(func, "__doc__") self.func = func def __get__(self, obj, cls): @@ -26,7 +30,12 @@ def __get__(self, obj, cls): def memReport(): for obj in gc.get_objects(): if torch.is_tensor(obj): - print(type(obj), obj.size(), obj.element_size() * obj.nelement() / 1024 / 1024, file=sys.stderr) + print( + type(obj), + obj.size(), + obj.element_size() * obj.nelement() / 1024 / 1024, + file=sys.stderr, + ) def cpuStats(): @@ -35,8 +44,8 @@ def cpuStats(): print(psutil.virtual_memory()) # physical memory usage pid = os.getpid() py = psutil.Process(pid) - memoryUse = py.memory_info()[0] / 2. ** 30 # memory use in GB...I think - print('memory GB:', memoryUse, file=sys.stderr) + memoryUse = py.memory_info()[0] / 2.0 ** 30 # memory use in GB...I think + print("memory GB:", memoryUse, file=sys.stderr) def update(d, u): diff --git a/neural-model/utils/vocab.py b/neural-model/utils/vocab.py index 376ad19..ce411ab 100644 --- a/neural-model/utils/vocab.py +++ b/neural-model/utils/vocab.py @@ -10,20 +10,18 @@ --freq-cutoff= frequency cutoff [default: 2] """ +import json +import pickle from collections import Counter from itertools import chain +import sentencepiece as spm import torch -import pickle from docopt import docopt -import json -import sentencepiece as spm - from utils.grammar import Grammar - -SAME_VARIABLE_TOKEN = '' -END_OF_VARIABLE_TOKEN = '' +SAME_VARIABLE_TOKEN = "" +END_OF_VARIABLE_TOKEN = "" PAD_ID = 0 @@ -36,16 +34,17 @@ def __init__(self, subtoken_model_path=None): self.subtoken_model = spm.SentencePieceProcessor() self.subtoken_model.load(subtoken_model_path) - vocab_path = subtoken_model_path[: subtoken_model_path.rfind('.model')] + '.vocab' + subtoken_model = subtoken_model_path[: subtoken_model_path.rfind(".model")] + vocab_path = f"{subtoken_model}.vocab" for i, line in enumerate(open(vocab_path)): - word, prob = line.strip().split('\t') + word, prob = line.strip().split("\t") self.word2id[word] = len(self.word2id) else: self.subtoken_model = None - self.word2id[''] = PAD_ID - self.word2id[''] = 1 - self.word2id[''] = 2 - self.word2id[''] = 3 + self.word2id[""] = PAD_ID + self.word2id[""] = 1 + self.word2id[""] = 2 + self.word2id[""] = 3 self.word2id[SAME_VARIABLE_TOKEN] = 4 self.id2word = {v: k for k, v in self.word2id.items()} @@ -55,7 +54,7 @@ def __getitem__(self, word): @property def unk_id(self): - return self.word2id[''] + return self.word2id[""] def is_unk(self, word): return word not in self.word2id @@ -64,13 +63,13 @@ def __contains__(self, word): return word in self.word2id def __setitem__(self, key, value): - raise ValueError('vocabulary is readonly') + raise ValueError("vocabulary is readonly") def __len__(self): return len(self.word2id) def __repr__(self): - return 'Vocabulary[size=%d]' % len(self) + return "Vocabulary[size=%d]" % len(self) def id2word(self, wid): return self.id2word[wid] @@ -85,34 +84,37 @@ def add(self, word): @property def params(self): - params = dict(unk_id=self.unk_id, word2id=self.word2id, - subtoken_model_path=self.subtoken_model_path) - if 'word_freq' in self.__dict__: - params['word_freq'] = self.word_freq + params = dict( + unk_id=self.unk_id, + word2id=self.word2id, + subtoken_model_path=self.subtoken_model_path, + ) + if "word_freq" in self.__dict__: + params["word_freq"] = self.word_freq return params def save(self, path): - json.dump(self.params, open(path, 'w'), indent=2) + json.dump(self.params, open(path, "w"), indent=2) @classmethod def load(cls, path=None, params=None): if path: - params = json.load(open(path, 'r')) + params = json.load(open(path, "r")) else: - assert params, 'Params must be given when path is None!' + assert params, "Params must be given when path is None!" - if 'subtoken_model_path' in params: - subtoken_model_path = params['subtoken_model_path'] + if "subtoken_model_path" in params: + subtoken_model_path = params["subtoken_model_path"] else: subtoken_model_path = None entry = VocabEntry(subtoken_model_path) - setattr(entry, 'word2id', params['word2id']) - setattr(entry, 'id2word', {v: k for k, v in params['word2id'].items()}) - if 'word_freq' in params: - setattr(entry, 'word_freq', params['word_freq']) + setattr(entry, "word2id", params["word2id"]) + setattr(entry, "id2word", {v: k for k, v in params["word2id"].items()}) + if "word_freq" in params: + setattr(entry, "word_freq", params["word_freq"]) return entry @@ -122,11 +124,14 @@ def from_corpus(corpus, size, freq_cutoff=0, predefined_tokens=None): word_freq = Counter(chain(*corpus)) freq_words = [w for w in word_freq if word_freq[w] >= freq_cutoff] - print('number of word types: %d, number of word types w/ frequency >= %d: %d' % (len(word_freq), freq_cutoff, - len(freq_words))) + print( + f"number of word types: {len(word_freq)}, " + f"number of word types w/ frequency >= {freq_cutoff}: " + f"{freq_words}" + ) top_k_words = sorted(word_freq, key=lambda x: (-word_freq[x], x))[:size] - print('top 30 words: %s' % ', '.join(top_k_words[:30])) + print("top 30 words: %s" % ", ".join(top_k_words[:30])) if predefined_tokens: for token in predefined_tokens: @@ -138,7 +143,7 @@ def from_corpus(corpus, size, freq_cutoff=0, predefined_tokens=None): vocab_entry.add(word) # store the work frequency table in the - setattr(vocab_entry, 'word_freq', word_freq) + setattr(vocab_entry, "word_freq", word_freq) return vocab_entry @@ -152,7 +157,10 @@ def __init__(self, **kwargs): self.entries.append(key) def __repr__(self): - return 'Vocab(%s)' % (', '.join('%s %swords' % (entry, getattr(self, entry)) for entry in self.entries)) + words = ", ".join( + f"{entry} {getattr(self, entry)}words" for entry in self.entries + ) + return f"Vocab({words})" @property def params(self): @@ -163,14 +171,14 @@ def params(self): return params def save(self, path): - json.dump(self.params, open(path, 'w'), indent=2) + json.dump(self.params, open(path, "w"), indent=2) @classmethod def load(cls, path): - params = json.load(open(path, 'r')) + params = json.load(open(path, "r")) entries = dict() for key, val in params.items(): - if key in ('grammar', ): + if key in ("grammar",): entry = Grammar.load(val) else: entry = VocabEntry.load(params=val) @@ -178,17 +186,17 @@ def load(cls, path): return cls(**entries) -if __name__ == '__main__': +if __name__ == "__main__": from utils.dataset import Dataset args = docopt(__doc__) - vocab_size = int(args['--size']) - vocab_file = args['VOCAB_FILE'] - train_set = Dataset(args['TRAIN_FILE']) + vocab_size = int(args["--size"]) + vocab_file = args["VOCAB_FILE"] + train_set = Dataset(args["TRAIN_FILE"]) - src_code_tokens_file = vocab_file + '.src_code_tokens.txt' + src_code_tokens_file = vocab_file + ".src_code_tokens.txt" src_preserved_tokens = set() - f_src_token = open(src_code_tokens_file, 'w') + f_src_token = open(src_code_tokens_file, "w") # extract vocab and node types node_types = set() @@ -209,65 +217,99 @@ def load(cls, path): if old_var_name != new_var_name: tgt_words.append(new_var_name) - if node.node_type == 'obj' or node.node_type == 'block' and hasattr(node, 'name'): + if ( + node.node_type == "obj" + or node.node_type == "block" + and hasattr(node, "name") + ): identifier_names.append(node.name) - if hasattr(node, 'type_tokens'): + if hasattr(node, "type_tokens"): type_tokens.extend(node.type_tokens) code_tokens = example.code_tokens - preserved_tokens = [token for token in code_tokens if token.startswith('@@') and token.endswith('@@')] + preserved_tokens = [ + token + for token in code_tokens + if token.startswith("@@") and token.endswith("@@") + ] src_preserved_tokens.update(preserved_tokens) - f_src_token.write(' '.join(code_tokens) + '\n') + f_src_token.write(" ".join(code_tokens) + "\n") f_src_token.close() - print('building source words vocabulary') - src_var_vocab_entry = VocabEntry.from_corpus([src_words], size=vocab_size, - freq_cutoff=int(args['--freq-cutoff'])) + print("building source words vocabulary") + src_var_vocab_entry = VocabEntry.from_corpus( + [src_words], size=vocab_size, freq_cutoff=int(args["--freq-cutoff"]) + ) - if args['--use-bpe']: - print('use bpe') + if args["--use-bpe"]: + print("use bpe") - print('building source code tokens vocabulary') + print("building source code tokens vocabulary") # train subtoken models - src_preserved_tokens = ','.join(src_preserved_tokens) - spm.SentencePieceTrainer.Train(f'--add_dummy_prefix=false --pad_id={PAD_ID} --bos_id=1 --eos_id=2 --unk_id=3 ' - f'--user_defined_symbols={src_preserved_tokens} ' - f'--vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.src_code_tokens --model_type=bpe ' - f'--input={src_code_tokens_file}') - src_code_tokens_vocab_entry = VocabEntry(vocab_file + '.src_code_tokens.model') - - print('building target words vocabulary') - tgt_word_file = args['VOCAB_FILE'] + '.var_names.txt' - with open(tgt_word_file, 'w') as f: + src_preserved_tokens = ",".join(src_preserved_tokens) + spm.SentencePieceTrainer.Train( + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + f"--user_defined_symbols={src_preserved_tokens} " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.src_code_tokens " + "--model_type=bpe " + f"--input={src_code_tokens_file}" + ) + src_code_tokens_vocab_entry = VocabEntry(f"{vocab_file}.src_code_tokens.model") + + print("building target words vocabulary") + tgt_word_file = args["VOCAB_FILE"] + ".var_names.txt" + with open(tgt_word_file, "w") as f: for name in tgt_words: - f.write(name + '\n') - - spm.SentencePieceTrainer.Train(f'--add_dummy_prefix=false --pad_id={PAD_ID} --bos_id=1 --eos_id=2 --unk_id=3 ' - f'--control_symbols= ' - f'--vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.tgt --model_type=bpe ' - f'--input={tgt_word_file}') - tgt_var_vocab_entry = VocabEntry(vocab_file + '.tgt.model') + f.write(name + "\n") + + spm.SentencePieceTrainer.Train( + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + "--control_symbols= " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.tgt " + "--model_type=bpe " + f"--input={tgt_word_file}" + ) + tgt_var_vocab_entry = VocabEntry(f"{vocab_file}.tgt.model") else: - tgt_var_vocab_entry = VocabEntry.from_corpus([tgt_words], size=vocab_size, - freq_cutoff=int(args['--freq-cutoff']), - predefined_tokens=[SAME_VARIABLE_TOKEN]) - - id_names_file = vocab_file + '.id_names.txt' - with open(id_names_file, 'w') as f: + tgt_var_vocab_entry = VocabEntry.from_corpus( + [tgt_words], + size=vocab_size, + freq_cutoff=int(args["--freq-cutoff"]), + predefined_tokens=[SAME_VARIABLE_TOKEN], + ) + + id_names_file = vocab_file + ".id_names.txt" + with open(id_names_file, "w") as f: for name in identifier_names: - f.write(name + '\n') + f.write(name + "\n") - print('train subtoken model for obj names') + print("train subtoken model for obj names") # train subtoken models - spm.SentencePieceTrainer.Train(f'--add_dummy_prefix=false --pad_id={PAD_ID} --bos_id=1 --eos_id=2 --unk_id=3 ' - f'--control_symbols= --vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.obj_name --model_type=bpe ' - f'--input={id_names_file}') - obj_name_vocab_entry = VocabEntry(vocab_file + '.obj_name.model') + spm.SentencePieceTrainer.Train( + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + "--control_symbols= " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.obj_name " + "--model_type=bpe " + f"--input={id_names_file}" + ) + obj_name_vocab_entry = VocabEntry(f"{vocab_file}.obj_name.model") type_vocab = Counter(type_tokens) num_types = 100 @@ -277,16 +319,18 @@ def load(cls, path): print(type_token, freq) var_types.append(type_token) - print('init node types and variable types') + print("init node types and variable types") grammar = Grammar(node_types, var_types) - print('Node types:', node_types) - print('Variable types:', var_types) + print("Node types:", node_types) + print("Variable types:", var_types) - vocab = Vocab(source=src_var_vocab_entry, - source_tokens=src_code_tokens_vocab_entry, - target=tgt_var_vocab_entry, - obj_name=obj_name_vocab_entry, - grammar=grammar) + vocab = Vocab( + source=src_var_vocab_entry, + source_tokens=src_code_tokens_vocab_entry, + target=tgt_var_vocab_entry, + obj_name=obj_name_vocab_entry, + grammar=grammar, + ) - vocab.save(args['VOCAB_FILE']) + vocab.save(args["VOCAB_FILE"]) diff --git a/prediction-plugin/model/attentional_recurrent_subtoken_decoder.py b/prediction-plugin/model/attentional_recurrent_subtoken_decoder.py index c6d4b18..42fd5b1 100644 --- a/prediction-plugin/model/attentional_recurrent_subtoken_decoder.py +++ b/prediction-plugin/model/attentional_recurrent_subtoken_decoder.py @@ -1,55 +1,48 @@ import sys -from collections import OrderedDict -from typing import List, Any +from collections import OrderedDict, namedtuple +from typing import Any, Dict, List import torch -from torch import nn as nn - +from model.decoder import Decoder from model.encoder import Encoder -from utils import util, nn_util -from utils.dataset import Example -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN, END_OF_VARIABLE_TOKEN - from model.recurrent_subtoken_decoder import RecurrentSubtokenDecoder +from torch import nn as nn +from utils import nn_util, util +from utils.ast import AbstractSyntaxTree +from utils.dataset import Example +from utils.vocab import END_OF_VARIABLE_TOKEN, SAME_VARIABLE_TOKEN, Vocab class AttentionalRecurrentSubtokenDecoder(RecurrentSubtokenDecoder): - def __init__(self, - variable_encoding_size: int, - context_encoding_size: int, - hidden_size: int, - dropout: float, - tie_embed: bool, - input_feed: bool, - vocab: Vocab): + def __init__( + self, + variable_encoding_size: int, + context_encoding_size: int, + hidden_size: int, + dropout: float, + tie_embed: bool, + input_feed: bool, + vocab: Vocab, + ): super(AttentionalRecurrentSubtokenDecoder, self).__init__( - variable_encoding_size, - hidden_size, - dropout, - tie_embed, - input_feed, - vocab + variable_encoding_size, hidden_size, dropout, tie_embed, input_feed, vocab ) - self.att_src_linear = nn.Linear( - context_encoding_size, - hidden_size, - bias=False - ) + self.att_src_linear = nn.Linear(context_encoding_size, hidden_size, bias=False) self.att_vec_linear = nn.Linear( - context_encoding_size + hidden_size, - hidden_size, - bias=False + context_encoding_size + hidden_size, hidden_size, bias=False ) @classmethod def default_params(cls): params = RecurrentSubtokenDecoder.default_params() - params.update({ - 'remove_duplicates_in_prediction': True, - 'context_encoding_size': 128, - 'attention_target': 'ast_nodes' # terminal_nodes - }) + params.update( + { + "remove_duplicates_in_prediction": True, + "context_encoding_size": 128, + "attention_target": "ast_nodes", # terminal_nodes + } + ) return params @@ -57,15 +50,15 @@ def default_params(cls): def build(cls, config): params = util.update(cls.default_params(), config) - vocab = Vocab.load(params['vocab_file']) + vocab = Vocab.load(params["vocab_file"]) model = cls( - params['variable_encoding_size'], - params['context_encoding_size'], - params['hidden_size'], - params['dropout'], - params['tie_embedding'], - params['input_feed'], - vocab + params["variable_encoding_size"], + params["context_encoding_size"], + params["hidden_size"], + params["dropout"], + params["tie_embedding"], + params["input_feed"], + vocab, ) model.config = params @@ -76,9 +69,9 @@ def rnn_step(self, x, h_tm1, context_encoding): ctx_t, alpha_t = nn_util.dot_prod_attention( h_t, - context_encoding['attention_value'], - context_encoding['attention_key'], - context_encoding['attention_value_mask'] + context_encoding["attention_value"], + context_encoding["attention_key"], + context_encoding["attention_value_mask"], ) att_t = torch.tanh(self.att_vec_linear(torch.cat([h_t, ctx_t], 1))) @@ -87,14 +80,15 @@ def rnn_step(self, x, h_tm1, context_encoding): return (h_t, cell_t), att_t, ctx_t def get_attention_memory(self, context_encoding): - att_tgt = self.config['attention_target'] - value, mask = \ - self.encoder.get_attention_memory(context_encoding, att_tgt) + att_tgt = self.config["attention_target"] + value, mask = self.encoder.get_attention_memory(context_encoding, att_tgt) key = self.att_src_linear(value) - return {'attention_key': key, - 'attention_value': value, - 'attention_value_mask': mask} + return { + "attention_key": key, + "attention_value": value, + "attention_value_mask": mask, + } def forward(self, src_ast_encoding, prediction_target): # prepare tensors for attention @@ -102,25 +96,22 @@ def forward(self, src_ast_encoding, prediction_target): src_ast_encoding.update(attention_memory) return RecurrentSubtokenDecoder.forward( - self, - src_ast_encoding, - prediction_target + self, src_ast_encoding, prediction_target ) def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: batch_size = len(examples) - beam_size = self.config['beam_size'] + beam_size = self.config["beam_size"] same_variable_id = self.vocab.target[SAME_VARIABLE_TOKEN] end_of_variable_id = self.vocab.target[END_OF_VARIABLE_TOKEN] - remove_duplicate = self.config['remove_duplicates_in_prediction'] + remove_duplicate = self.config["remove_duplicates_in_prediction"] variable_nums = [] for ast_id, example in enumerate(examples): variable_nums.append(len(example.ast.variables)) beams = OrderedDict( - (ast_id, [self.Hypothesis([[]], 0, 0.)]) - for ast_id in range(batch_size) + (ast_id, [self.Hypothesis([[]], 0, 0.0)]) for ast_id in range(batch_size) ) hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) completed_hyps = [[] for _ in range(batch_size)] @@ -134,47 +125,40 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: attention_memory = self.get_attention_memory(context_encoding) context_encoding_t = attention_memory - h_tm1 = self.get_init_state(context_encoding) + h_tm1 = _h_0 = self.get_init_state(context_encoding) # Note that we are using the `restoration_indices` from # `context_encoding`, which is the word-level restoration index # (batch_size, variable_master_node_num, encoding_size) - variable_encoding = context_encoding['variable_encoding'] + variable_encoding = context_encoding["variable_encoding"] # (batch_size, encoding_size) variable_name_embed_tm1 = att_tm1 = torch.zeros( - batch_size, - self.lstm_cell.hidden_size, - device=self.device + batch_size, self.lstm_cell.hidden_size, device=self.device ) - max_prediction_time_step = self.config['max_prediction_time_step'] + max_prediction_time_step = self.config["max_prediction_time_step"] for t in range(0, max_prediction_time_step): # (total_live_hyp_num, encoding_size) if t > 0: - variable_encoding_t = \ - variable_encoding[hyp_ast_ids_t, hyp_variable_ptrs_t] + variable_encoding_t = variable_encoding[ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] else: variable_encoding_t = variable_encoding[:, 0] - if self.config['input_feed']: + if self.config["input_feed"]: x = torch.cat( - [variable_encoding_t, variable_name_embed_tm1, att_tm1], - dim=-1 + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 ) else: - x = torch.cat( - [variable_encoding_t, variable_name_embed_tm1], - dim=-1 - ) + x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) h_t, q_t, alpha_t = self.rnn_step(x, h_tm1, context_encoding_t) # (total_live_hyp_num, vocab_size) - hyp_var_name_scores_t = \ - torch.log_softmax(self.state2names(q_t), dim=-1) + hyp_var_name_scores_t = torch.log_softmax(self.state2names(q_t), dim=-1) - cont_cand_hyp_scores = \ - hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t + cont_cand_hyp_scores = hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t new_beams = OrderedDict() live_beam_ids = [] @@ -188,13 +172,13 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: for beam_id, (ast_id, beam) in enumerate(beams.items()): beam_end_hyp_pos = beam_start_hyp_pos + len(beam) # (live_beam_size, vocab_size) - beam_cont_cand_hyp_scores = \ - cont_cand_hyp_scores[beam_start_hyp_pos: beam_end_hyp_pos] + beam_cont_cand_hyp_scores = cont_cand_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] cont_beam_size = beam_size - len(completed_hyps[ast_id]) - beam_new_hyp_scores, beam_new_hyp_positions = \ - torch.topk(beam_cont_cand_hyp_scores.view(-1), - k=cont_beam_size, - dim=-1) + beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( + beam_cont_cand_hyp_scores.view(-1), k=cont_beam_size, dim=-1 + ) # (cont_beam_size) beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size @@ -212,8 +196,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: variable_ptr = prev_hyp.variable_ptr new_variable_list = list(prev_hyp.variable_list) - new_variable_list[-1] = \ - list(new_variable_list[-1] + [hyp_var_name_id]) + new_variable_list[-1] = list( + new_variable_list[-1] + [hyp_var_name_id] + ) if hyp_var_name_id == end_of_variable_id: # remove empty cases @@ -222,18 +207,21 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if remove_duplicate: last_pred = new_variable_list[-1] - if any(x == last_pred - for x in new_variable_list[:-1] - if x != [same_variable_id, end_of_variable_id] + if any( + x == last_pred + for x in new_variable_list[:-1] + if x != [same_variable_id, end_of_variable_id] ): continue variable_ptr += 1 new_variable_list.append([]) - new_hyp = self.Hypothesis(variable_list=new_variable_list, - variable_ptr=variable_ptr, - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=new_variable_list, + variable_ptr=variable_ptr, + score=new_hyp_score, + ) if variable_ptr == variable_nums[ast_id]: completed_hyps[ast_id].append(new_hyp) @@ -245,7 +233,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: new_hyp_var_name_ids.append(hyp_var_name_id) new_hyp_ast_ids.append(ast_id) new_hyp_variable_ptrs.append(variable_ptr) - is_same_variable_mask.append(1. if prev_hyp.variable_ptr == variable_ptr else 0.) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) beam_start_hyp_pos = beam_end_hyp_pos @@ -254,10 +244,14 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: h_tm1 = (h_t[0][live_prev_hyp_ids], h_t[1][live_prev_hyp_ids]) att_tm1 = q_t[live_prev_hyp_ids] - if self.config['tie_embedding']: - variable_name_embed_tm1 = self.state2names.weight[new_hyp_var_name_ids] + if self.config["tie_embedding"]: + variable_name_embed_tm1 = self.state2names.weight[ + new_hyp_var_name_ids + ] else: - variable_name_embed_tm1 = self.var_name_embed.weight[new_hyp_var_name_ids] + variable_name_embed_tm1 = self.var_name_embed.weight[ + new_hyp_var_name_ids + ] hyp_ast_ids_t = new_hyp_ast_ids hyp_variable_ptrs_t = new_hyp_variable_ptrs @@ -265,15 +259,26 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: beams = new_beams # (total_hyp_num, max_tree_size, node_encoding_size) - context_encoding_t = dict(attention_key=attention_memory['attention_key'][hyp_ast_ids_t], - attention_value=attention_memory['attention_value'][hyp_ast_ids_t], - attention_value_mask=attention_memory['attention_value_mask'][hyp_ast_ids_t]) + context_encoding_t = dict( + attention_key=attention_memory["attention_key"][hyp_ast_ids_t], + attention_value=attention_memory["attention_value"][hyp_ast_ids_t], + attention_value_mask=attention_memory["attention_value_mask"][ + hyp_ast_ids_t + ], + ) if self.independent_prediction_for_each_variable: - is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) - h_tm1 = (h_tm1[0] * is_same_variable_mask, h_tm1[1] * is_same_variable_mask) + is_same_variable_mask = torch.tensor( + is_same_variable_mask, device=self.device, dtype=torch.float + ).unsqueeze(-1) + h_tm1 = ( + h_tm1[0] * is_same_variable_mask, + h_tm1[1] * is_same_variable_mask, + ) att_tm1 = att_tm1 * is_same_variable_mask - variable_name_embed_tm1 = variable_name_embed_tm1 * is_same_variable_mask + variable_name_embed_tm1 = ( + variable_name_embed_tm1 * is_same_variable_mask + ) else: break @@ -284,15 +289,21 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if not hyps: # return identity renamings - print(f'Failed to found a hypothesis for function {ast.compilation_unit}', file=sys.stderr) + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) variable_rename_result = dict() for old_name in ast.variables: - variable_rename_result[old_name] = {'new_name': old_name, - 'prob': 0.} + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } example_rename_results = [variable_rename_result] else: example_rename_results = [] + for hyp in hyps: variable_rename_result = dict() for var_id, old_name in enumerate(ast.variables): @@ -300,9 +311,14 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Any]: if var_name_token_ids == [same_variable_id, end_of_variable_id]: new_var_name = old_name else: - new_var_name = self.vocab.target.subtoken_model.decode_ids(var_name_token_ids) - - variable_rename_result[old_name] = {'new_name': new_var_name, 'prob': hyp.score} + new_var_name = self.vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": hyp.score, + } example_rename_results.append(variable_rename_result) diff --git a/prediction-plugin/model/embedding.py b/prediction-plugin/model/embedding.py index cd74c60..3577d69 100644 --- a/prediction-plugin/model/embedding.py +++ b/prediction-plugin/model/embedding.py @@ -1,10 +1,8 @@ -from typing import List, Dict +from typing import Dict, List import sentencepiece as spm - import torch import torch.nn as nn - from utils.nn_util import SMALL_NUMBER @@ -16,23 +14,25 @@ def __init__(self, bpe_model_path, embedding_size): self.bpe_model.load(bpe_model_path) self.size = len(self.bpe_model) self.pad_id = self.bpe_model.pad_id() - self.embeddings = nn.Embedding(self.size, - embedding_size, - padding_idx=self.pad_id) + self.embeddings = nn.Embedding( + self.size, embedding_size, padding_idx=self.pad_id + ) @property def device(self): return self.embeddings.weight.device @classmethod - def to_input_tensor(cls, - sub_tokens_list: List[List[str]], - bpe_model: spm.SentencePieceProcessor, - pad_id=0) -> torch.Tensor: + def to_input_tensor( + cls, + sub_tokens_list: List[List[str]], + bpe_model: spm.SentencePieceProcessor, + pad_id=0, + ) -> torch.Tensor: max_subword_num = max(len(x) for x in sub_tokens_list) - idx_tensor = torch.zeros(len(sub_tokens_list), - max_subword_num, - dtype=torch.long) + idx_tensor = torch.zeros( + len(sub_tokens_list), max_subword_num, dtype=torch.long + ) idx_tensor.fill_(pad_id) for i, token_list in enumerate(sub_tokens_list): @@ -42,8 +42,7 @@ def to_input_tensor(cls, return idx_tensor def get_embedding(self, sub_tokens_list: List[List[str]]) -> torch.Tensor: - idx_tensor = \ - self.to_input_tensor(sub_tokens_list, self.bpe_model, self.pad_id) + idx_tensor = self.to_input_tensor(sub_tokens_list, self.bpe_model, self.pad_id) embedding = self.forward(idx_tensor) return embedding @@ -51,10 +50,9 @@ def get_embedding(self, sub_tokens_list: List[List[str]]) -> torch.Tensor: def forward(self, sub_tokens_indices: torch.Tensor) -> torch.Tensor: # sub_tokens_indices: (batch_size, max_sub_token_num) sub_tokens_mask = torch.ne(sub_tokens_indices, self.pad_id).float() - embedding = \ - self.embeddings(sub_tokens_indices) * sub_tokens_mask.unsqueeze(-1) - embedding = \ - embedding.sum(dim=1) / sub_tokens_mask.sum(-1).unsqueeze(-1) + embedding = self.embeddings(sub_tokens_indices) * sub_tokens_mask.unsqueeze(-1) + embedding = embedding.sum(dim=1) / sub_tokens_mask.sum(-1).unsqueeze(-1) + return embedding @@ -64,36 +62,37 @@ def __init__(self, type_num, embedding_size, pad_id=0): self.size = type_num self.pad_id = pad_id - self.embeddings = \ - nn.Embedding(type_num, embedding_size, padding_idx=self.pad_id) + self.embeddings = nn.Embedding( + type_num, embedding_size, padding_idx=self.pad_id + ) @property def device(self): return self.embeddings.weight.device @classmethod - def to_input_tensor(cls, token_types_list: List[List[str]], - type2id: Dict[str, int], - pad_id=0) -> torch.Tensor: + def to_input_tensor( + cls, token_types_list: List[List[str]], type2id: Dict[str, int], pad_id=0 + ) -> torch.Tensor: max_subword_num = max(len(x) for x in token_types_list) - idx_tensor = torch.zeros(len(token_types_list), - max_subword_num, - dtype=torch.long) + idx_tensor = torch.zeros( + len(token_types_list), max_subword_num, dtype=torch.long + ) idx_tensor.fill_(pad_id) for i, token_list in enumerate(token_types_list): - idx_tensor[i, :len(token_list)] = \ - torch.tensor([type2id[t] for t in token_list]) + idx_tensor[i, : len(token_list)] = torch.tensor( + [type2id[t] for t in token_list] + ) + return idx_tensor def forward(self, type_tokens_indices: torch.Tensor) -> torch.Tensor: # type_tokens_indices: (batch_size, max_sub_token_num) sub_tokens_mask = torch.ne(type_tokens_indices, self.pad_id).float() - embedding = \ - self.embeddings(type_tokens_indices) \ - * sub_tokens_mask.unsqueeze(-1) - embedding = \ - embedding.sum(dim=1) \ - / (sub_tokens_mask.sum(-1) + SMALL_NUMBER).unsqueeze(-1) + embedding = self.embeddings(type_tokens_indices) * sub_tokens_mask.unsqueeze(-1) + embedding = embedding.sum(dim=1) / ( + sub_tokens_mask.sum(-1) + SMALL_NUMBER + ).unsqueeze(-1) return embedding diff --git a/prediction-plugin/model/ensemble_model.py b/prediction-plugin/model/ensemble_model.py new file mode 100644 index 0000000..c915adc --- /dev/null +++ b/prediction-plugin/model/ensemble_model.py @@ -0,0 +1,319 @@ +from collections import OrderedDict, namedtuple + +from model.model import * +from utils.vocab import END_OF_VARIABLE_TOKEN + + +class EnsembleModel(nn.Module): + def __init__(self, model_paths, use_cuda, config): + super(EnsembleModel, self).__init__() + + self.models = nn.ModuleList() + for model_path in model_paths: + model = RenamingModel.load(model_path, use_cuda=use_cuda) + self.models.append(model) + + self.max_prediction_time_step = min( + model.decoder.config["max_prediction_time_step"] for model in self.models + ) + self.device = torch.device("cuda:0") if use_cuda else torch.device("cpu") + + self.Hypothesis = namedtuple( + "Hypothesis", ["variable_list", "variable_ptr", "score"] + ) + self.config = config + print("Ensemble Model Config", file=sys.stderr) + print(self.config, file=sys.stderr) + + @classmethod + def default_params(cls): + return { + "encoder": {"type": "EnsembleModel"}, + "decoder": {"beam_size": 5, "remove_duplicates_in_prediction": False}, + "train": { + "eval_batch_size": 100, + "num_batchers": 5, + "num_readers": 5, + "buffer_size": 100, + }, + } + + @classmethod + def load(cls, model_paths, use_cuda=False, new_config=None): + params = util.update(cls.default_params(), new_config) + model = cls(model_paths, use_cuda, params) + + return model + + def predict(self, examples): + context_encodings = [] + attention_memories = [] + variable_encodings = [] + remove_duplicate = self.config["decoder"]["remove_duplicates_in_prediction"] + + h_tm1 = [] + variable_name_embed_tm1 = [] + + for model in self.models: + tensor_dict = model.batcher.to_tensor_dict(examples) + nn_util.to(tensor_dict, self.device) + context_encoding = model.encoder(tensor_dict) + # prepare tensors for attention + attention_memory = model.decoder.get_attention_memory(context_encoding) + h_0 = model.decoder.get_init_state(context_encoding) + + context_encodings.append(context_encoding) + variable_encodings.append(context_encoding["variable_encoding"]) + attention_memories.append(attention_memory) + h_tm1.append(h_0) + + att_tm1 = torch.zeros( + len(examples), model.decoder.lstm_cell.hidden_size, device=self.device + ) + variable_name_embed_tm1.append(att_tm1) + + context_encoding_t = attention_memories + + beam_size = self.config["decoder"]["beam_size"] + batch_size = len(examples) + vocab = self.models[0].decoder.vocab + same_variable_id = vocab.target[SAME_VARIABLE_TOKEN] + end_of_variable_id = vocab.target[END_OF_VARIABLE_TOKEN] + variable_nums = [len(e.ast.variables) for e in examples] + + beams = OrderedDict( + (ast_id, [self.Hypothesis([[]], 0, 0.0)]) for ast_id in range(batch_size) + ) + hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) + completed_hyps = [[] for _ in range(batch_size)] + tgt_vocab_size = len(self.models[0].decoder.vocab.target) + + for t in range(0, self.max_prediction_time_step): + candidate_cont_hyp_scores: List[torch.Tensor] = [] + model_h_t = [] + model_q_t = [] + for model_id, model in enumerate(self.models): + # (total_live_hyp_num, encoding_size) + if t > 0: + variable_encoding_t = variable_encodings[model_id][ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] + else: + variable_encoding_t = variable_encodings[model_id][:, 0] + + if model.decoder.config["input_feed"]: + x = torch.cat( + [ + variable_encoding_t, + variable_name_embed_tm1[model_id], + att_tm1, + ], + dim=-1, + ) + else: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1[model_id]], dim=-1 + ) + + h_t, q_t, alpha_t = model.decoder.rnn_step( + x, h_tm1[model_id], context_encoding_t[model_id] + ) + model_h_t.append(h_t) + model_q_t.append(q_t) + + # (total_live_hyp_num, vocab_size) + hyp_var_name_scores_t = torch.log_softmax( + model.decoder.state2names(q_t), dim=-1 + ) + cont_cand_hyp_scores = ( + hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t + ) + candidate_cont_hyp_scores.append(cont_cand_hyp_scores) + + # (total_live_hyp_num, vocab_size) + candidate_cont_hyp_scores = torch.logsumexp( + torch.stack(candidate_cont_hyp_scores, dim=-1), dim=-1 + ) + + new_beams = OrderedDict() + live_beam_ids = [] + new_hyp_scores = [] + live_prev_hyp_ids = [] + new_hyp_var_name_ids = [] + new_hyp_ast_ids = [] + new_hyp_variable_ptrs = [] + is_same_variable_mask = [] + beam_start_hyp_pos = 0 + + for beam_id, (ast_id, beam) in enumerate(beams.items()): + beam_end_hyp_pos = beam_start_hyp_pos + len(beam) + # (live_beam_size, vocab_size) + beam_cont_cand_hyp_scores = candidate_cont_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] + + cont_beam_size = beam_size - len(completed_hyps[ast_id]) + beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( + beam_cont_cand_hyp_scores.view(-1), k=cont_beam_size, dim=-1 + ) + + # (cont_beam_size) + beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size + beam_hyp_var_name_ids = beam_new_hyp_positions % tgt_vocab_size + + _prev_hyp_ids = beam_prev_hyp_ids.cpu() + _hyp_var_name_ids = beam_hyp_var_name_ids.cpu() + _new_hyp_scores = beam_new_hyp_scores.cpu() + + for i in range(cont_beam_size): + prev_hyp_id = _prev_hyp_ids[i].item() + prev_hyp = beam[prev_hyp_id] + hyp_var_name_id = _hyp_var_name_ids[i].item() + new_hyp_score = _new_hyp_scores[i].item() + + variable_ptr = prev_hyp.variable_ptr + new_variable_list = list(prev_hyp.variable_list) + new_variable_list[-1] = list( + new_variable_list[-1] + [hyp_var_name_id] + ) + + if hyp_var_name_id == end_of_variable_id: + # remove empty cases + if new_variable_list[-1] == [end_of_variable_id]: + continue + + if remove_duplicate: + last_pred = new_variable_list[-1] + if any( + x == last_pred + for x in new_variable_list[:-1] + if x != [same_variable_id, end_of_variable_id] + ): + # print('found a duplicate!', ', '.join([str(x) for x in last_pred])) + continue + + variable_ptr += 1 + new_variable_list.append([]) + + new_hyp = self.Hypothesis( + variable_list=new_variable_list, + variable_ptr=variable_ptr, + score=new_hyp_score, + ) + + if variable_ptr == variable_nums[ast_id]: + completed_hyps[ast_id].append(new_hyp) + else: + new_beams.setdefault(ast_id, []).append(new_hyp) + live_beam_ids.append(beam_id) + new_hyp_scores.append(new_hyp_score) + live_prev_hyp_ids.append(beam_start_hyp_pos + prev_hyp_id) + new_hyp_var_name_ids.append(hyp_var_name_id) + new_hyp_ast_ids.append(ast_id) + new_hyp_variable_ptrs.append(variable_ptr) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) + + beam_start_hyp_pos = beam_end_hyp_pos + + if live_beam_ids: + hyp_scores_tm1 = torch.tensor(new_hyp_scores, device=self.device) + h_tm1 = [ + ( + model_h_t[model_id][0][live_prev_hyp_ids], + model_h_t[model_id][1][live_prev_hyp_ids], + ) + for model_id in range(len(self.models)) + ] + att_tm1 = [ + model_q_t[model_id][live_prev_hyp_ids] + for model_id in range(len(self.models)) + ] + + variable_name_embed_tm1 = [ + model.decoder.state2names.weight[new_hyp_var_name_ids] + for model in self.models + ] + hyp_ast_ids_t = new_hyp_ast_ids + hyp_variable_ptrs_t = new_hyp_variable_ptrs + + beams = new_beams + + # (total_hyp_num, max_tree_size, node_encoding_size) + context_encoding_t = [ + dict( + attention_key=attention_memories[model_id]["attention_key"][ + hyp_ast_ids_t + ], + attention_value=attention_memories[model_id]["attention_value"][ + hyp_ast_ids_t + ], + attention_value_mask=attention_memories[model_id][ + "attention_value_mask" + ][hyp_ast_ids_t], + ) + for model_id in range(len(self.models)) + ] + + # if self.independent_prediction_for_each_variable: + # is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) + # h_tm1 = [(h[0] * is_same_variable_mask, h[1] * is_same_variable_mask) for h in h_tm1] + # att_tm1 = [x * is_same_variable_mask for x in att_tm1] + # variable_name_embed_tm1 = [x * is_same_variable_mask for x in variable_name_embed_tm1] + else: + break + + variable_rename_results = [] + for i, hyps in enumerate(completed_hyps): + variable_rename_result = dict() + ast = examples[i].ast + hyps = sorted(hyps, key=lambda hyp: -hyp.score) + + if not hyps: + # return identity renamings + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) + for old_name in ast.variables: + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } + else: + top_hyp = hyps[0] + # sub_token_ptr = 0 + # for old_name in ast.variables: + # sub_token_begin = sub_token_ptr + # while top_hyp.variable_list[sub_token_ptr] != end_of_variable_id: + # sub_token_ptr += 1 + # sub_token_ptr += 1 # point to first sub-token of next variable + # sub_token_end = sub_token_ptr + # + # var_name_token_ids = top_hyp.variable_list[sub_token_begin: sub_token_end] # include ending + # if var_name_token_ids == [same_variable_id, end_of_variable_id]: + # new_var_name = old_name + # else: + # new_var_name = vocab.target.subtoken_model.decode_ids(var_name_token_ids) + # + # variable_rename_result[old_name] = {'new_name': new_var_name, + # 'prob': top_hyp.score} + + for var_id, old_name in enumerate(ast.variables): + var_name_token_ids = top_hyp.variable_list[var_id] + if var_name_token_ids == [same_variable_id, end_of_variable_id]: + new_var_name = old_name + else: + new_var_name = vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } + + variable_rename_results.append(variable_rename_result) + + return variable_rename_results diff --git a/prediction-plugin/model/gnn.py b/prediction-plugin/model/gnn.py index ddbbdd6..8f6766c 100644 --- a/prediction-plugin/model/gnn.py +++ b/prediction-plugin/model/gnn.py @@ -6,13 +6,11 @@ class AdjacencyList: """represent the topology of a graph""" - def __init__(self, - node_num: int, - adj_list: List, - device: torch.device = None): + + def __init__(self, node_num: int, adj_list: List, device: torch.device = None): self.node_num = node_num if device is None: - device = torch.device('cpu') + device = torch.device("cpu") self.data = torch.tensor(adj_list, dtype=torch.long, device=device) self.edge_num = len(adj_list) @@ -20,7 +18,7 @@ def __init__(self, def device(self): return self.data.device - def to(self, device: torch.device) -> 'AdjacencyList': + def to(self, device: torch.device) -> "AdjacencyList": if self.data.device != device: self.data = self.data.to(device) return self @@ -30,19 +28,22 @@ def __getitem__(self, item): class GatedGraphNeuralNetwork(nn.Module): - def __init__(self, hidden_size, num_edge_types, layer_timesteps, - residual_connections, - state_to_message_dropout=0.3, - rnn_dropout=0.3, - use_bias_for_message_linear=True): + def __init__( + self, + hidden_size, + num_edge_types, + layer_timesteps, + residual_connections, + state_to_message_dropout=0.3, + rnn_dropout=0.3, + use_bias_for_message_linear=True, + ): super(GatedGraphNeuralNetwork, self).__init__() self.hidden_size = hidden_size self.num_edge_types = num_edge_types self.layer_timesteps = layer_timesteps - self.residual_connections = { - int(k): v for k, v in residual_connections.items() - } + self.residual_connections = {int(k): v for k, v in residual_connections.items()} self.state_to_message_dropout = state_to_message_dropout self.rnn_dropout = rnn_dropout self.use_bias_for_message_linear = use_bias_for_message_linear @@ -55,46 +56,48 @@ def __init__(self, hidden_size, num_edge_types, layer_timesteps, for layer_idx in range(len(self.layer_timesteps)): # Initiate a linear transformation for each edge type for edge_type_j in range(self.num_edge_types): - state_to_msg_linear_layer_i_type_j = \ - nn.Linear(self.hidden_size, - self.hidden_size, - bias=use_bias_for_message_linear) - layer_name = \ - f'state_to_message_linear_layer{layer_idx}' \ - f'_type{edge_type_j}' - self.state_to_message_linears[layer_name] = \ - state_to_msg_linear_layer_i_type_j - - layer_residual_connections = \ - self.residual_connections.get(layer_idx, []) - rnn_cell_layer_i = \ - nn.GRUCell(self.hidden_size - * (1 + len(layer_residual_connections)), - self.hidden_size) + state_to_msg_linear_layer_i_type_j = nn.Linear( + self.hidden_size, self.hidden_size, bias=use_bias_for_message_linear + ) + layer_name = ( + f"state_to_message_linear_layer{layer_idx}" f"_type{edge_type_j}" + ) + self.state_to_message_linears[ + layer_name + ] = state_to_msg_linear_layer_i_type_j + + layer_residual_connections = self.residual_connections.get(layer_idx, []) + rnn_cell_layer_i = nn.GRUCell( + self.hidden_size * (1 + len(layer_residual_connections)), + self.hidden_size, + ) self.rnn_cells.append(rnn_cell_layer_i) - self.state_to_message_dropout_layer = \ - nn.Dropout(self.state_to_message_dropout) + self.state_to_message_dropout_layer = nn.Dropout(self.state_to_message_dropout) self.rnn_dropout_layer = nn.Dropout(self.rnn_dropout) @property def device(self): return self.rnn_cells[0].weight_hh.device - def forward(self, - initial_node_representation: torch.Tensor, - adjacency_lists: List[AdjacencyList], - return_all_states=False) -> torch.Tensor: + def forward( + self, + initial_node_representation: torch.Tensor, + adjacency_lists: List[AdjacencyList], + return_all_states=False, + ) -> torch.Tensor: return self.compute_node_representations( initial_node_representation, adjacency_lists, - return_all_states=return_all_states + return_all_states=return_all_states, ) - def compute_node_representations(self, - initial_node_representation: torch.Tensor, - adjacency_lists: List[AdjacencyList], - return_all_states=False) -> torch.Tensor: + def compute_node_representations( + self, + initial_node_representation: torch.Tensor, + adjacency_lists: List[AdjacencyList], + return_all_states=False, + ) -> torch.Tensor: """ V: number of nodes on the batched graph E: number of edges on the batched graph @@ -107,20 +110,19 @@ def compute_node_representations(self, init_node_repr_size = initial_node_representation.size(1) if init_node_repr_size < self.hidden_size: pad_size = self.hidden_size - init_node_repr_size - zero_pads = torch.zeros(node_num, - pad_size, - dtype=torch.float, - device=self.device) - initial_node_representation = \ - torch.cat([initial_node_representation, zero_pads], dim=-1) + zero_pads = torch.zeros( + node_num, pad_size, dtype=torch.float, device=self.device + ) + initial_node_representation = torch.cat( + [initial_node_representation, zero_pads], dim=-1 + ) # Shape: [V, D] node_states_per_layer = [initial_node_representation] # list of tensors of message targets of shape [E] message_targets = [] - for edge_type_idx, adjacency_list_for_edge_type in \ - enumerate(adjacency_lists): + for edge_type_idx, adjacency_list_for_edge_type in enumerate(adjacency_lists): if adjacency_list_for_edge_type.edge_num > 0: edge_targets = adjacency_list_for_edge_type[:, 1] message_targets.append(edge_targets) @@ -134,12 +136,12 @@ def compute_node_representations(self, # M ~ number of messages (sum of all E) # Extract residual messages, if any: - layer_residual_connections = \ - self.residual_connections.get(layer_idx, []) + layer_residual_connections = self.residual_connections.get(layer_idx, []) # List[(V, D)] - layer_residual_states: List[torch.Tensor] = \ - [node_states_per_layer[residual_layer_idx] - for residual_layer_idx in layer_residual_connections] + layer_residual_states: List[torch.Tensor] = [ + node_states_per_layer[residual_layer_idx] + for residual_layer_idx in layer_residual_connections + ] # Record new states for this layer. Initialised to last state, but # will be updated below: @@ -152,26 +154,25 @@ def compute_node_representations(self, message_source_states: List[torch.Tensor] = [] # Collect incoming messages per edge type - for edge_type_idx, adjacency_list_for_edge_type in \ - enumerate(adjacency_lists): + for edge_type_idx, adjacency_list_for_edge_type in enumerate( + adjacency_lists + ): if adjacency_list_for_edge_type.edge_num > 0: # shape [E] edge_sources = adjacency_list_for_edge_type[:, 0] # shape [E, D] - edge_source_states = \ - node_states_for_this_layer[edge_sources] + edge_source_states = node_states_for_this_layer[edge_sources] - layer_name = \ - f'state_to_message_linear_layer{layer_idx}'\ - f'_type{edge_type_idx}' - f_state_to_message = \ - self.state_to_message_linears[layer_name] + layer_name = ( + f"state_to_message_linear_layer{layer_idx}" + f"_type{edge_type_idx}" + ) + f_state_to_message = self.state_to_message_linears[layer_name] # Shape [E, D] - all_messages_for_edge_type = \ - self.state_to_message_dropout_layer( - f_state_to_message(edge_source_states) - ) + all_messages_for_edge_type = self.state_to_message_dropout_layer( + f_state_to_message(edge_source_states) + ) messages.append(all_messages_for_edge_type) message_source_states.append(edge_source_states) @@ -181,35 +182,29 @@ def compute_node_representations(self, # Sum up messages that go to the same target node # shape [V, D] - incoming_messages = torch.zeros(node_num, - messages.size(1), - device=self.device) - incoming_messages = \ - incoming_messages.scatter_add_(0, - message_targets - .unsqueeze(-1) - .expand_as(messages), - messages) - node_target_num = \ - torch.zeros(node_num, device=self.device).scatter_add( - 0, - message_targets, - torch.ones(messages.size(0), device=self.device) - ) - incoming_messages = \ - incoming_messages / (node_target_num + 1e-7).unsqueeze(-1) + incoming_messages = torch.zeros( + node_num, messages.size(1), device=self.device + ) + incoming_messages = incoming_messages.scatter_add_( + 0, message_targets.unsqueeze(-1).expand_as(messages), messages + ) + node_target_num = torch.zeros(node_num, device=self.device).scatter_add( + 0, message_targets, torch.ones(messages.size(0), device=self.device) + ) + incoming_messages = incoming_messages / ( + node_target_num + 1e-7 + ).unsqueeze(-1) # shape [V, D * (1 + num of residual connections)] - incoming_information = \ - torch.cat(layer_residual_states + [incoming_messages], - dim=-1) + incoming_information = torch.cat( + layer_residual_states + [incoming_messages], dim=-1 + ) # pass updated vertex features into RNN cell # Shape [V, D] updated_node_states = self.rnn_cells[layer_idx]( incoming_information, node_states_for_this_layer ) - updated_node_states = \ - self.rnn_dropout_layer(updated_node_states) + updated_node_states = self.rnn_dropout_layer(updated_node_states) node_states_for_this_layer = updated_node_states node_states_per_layer.append(node_states_for_this_layer) @@ -222,25 +217,27 @@ def compute_node_representations(self, def main(): - gnn = GatedGraphNeuralNetwork(hidden_size=64, - num_edge_types=2, - layer_timesteps=[3, 5, 7, 2], - residual_connections={2: [0], 3: [0, 1]}) - - adj_list_type1 = AdjacencyList(node_num=4, - adj_list=[(0, 2), (2, 1), (1, 3)], - device=gnn.device) - adj_list_type2 = AdjacencyList(node_num=4, - adj_list=[(0, 0), (0, 1)], - device=gnn.device) + gnn = GatedGraphNeuralNetwork( + hidden_size=64, + num_edge_types=2, + layer_timesteps=[3, 5, 7, 2], + residual_connections={2: [0], 3: [0, 1]}, + ) + + adj_list_type1 = AdjacencyList( + node_num=4, adj_list=[(0, 2), (2, 1), (1, 3)], device=gnn.device + ) + adj_list_type2 = AdjacencyList( + node_num=4, adj_list=[(0, 0), (0, 1)], device=gnn.device + ) node_representations = gnn.compute_node_representations( initial_node_representation=torch.randn(4, 64), - adjacency_lists=[adj_list_type1, adj_list_type2] + adjacency_lists=[adj_list_type1, adj_list_type2], ) print(node_representations) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/prediction-plugin/model/graph_encoder.py b/prediction-plugin/model/graph_encoder.py index 1d68c02..3eaa689 100644 --- a/prediction-plugin/model/graph_encoder.py +++ b/prediction-plugin/model/graph_encoder.py @@ -1,13 +1,13 @@ -from typing import List, Dict, Tuple +from operator import itemgetter +from typing import Dict, List, Tuple import numpy as np import torch -from torch import nn as nn - from model.embedding import NodeTypeEmbedder, SubTokenEmbedder from model.encoder import Encoder -from model.gnn import GatedGraphNeuralNetwork, AdjacencyList +from model.gnn import AdjacencyList, GatedGraphNeuralNetwork from model.sequential_encoder import SequentialEncoder +from torch import nn as nn from utils import util from utils.ast import AbstractSyntaxTree from utils.grammar import Grammar @@ -17,15 +17,18 @@ class GraphASTEncoder(Encoder): """An encoder based on gated recurrent graph neural networks""" - def __init__(self, - gnn: GatedGraphNeuralNetwork, - connections: List[str], - node_syntax_type_embedding_size: int, - decoder_hidden_size: int, - node_type_embedder: NodeTypeEmbedder, - node_content_embedder: SubTokenEmbedder, - vocab: Vocab, - config): + + def __init__( + self, + gnn: GatedGraphNeuralNetwork, + connections: List[str], + node_syntax_type_embedding_size: int, + decoder_hidden_size: int, + node_type_embedder: NodeTypeEmbedder, + node_content_embedder: SubTokenEmbedder, + vocab: Vocab, + config, + ): super(GraphASTEncoder, self).__init__() self.connections = connections @@ -35,14 +38,14 @@ def __init__(self, self.grammar = grammar = vocab.grammar self.node_syntax_type_embedding = nn.Embedding( - len(grammar.syntax_types) + 2, - node_syntax_type_embedding_size + len(grammar.syntax_types) + 2, node_syntax_type_embedding_size ) self.variable_master_node_type_idx = len(grammar.syntax_types) self.master_node_type_idx = self.variable_master_node_type_idx + 1 - self.var_node_name_embedding = \ - nn.Embedding(len(vocab.source), gnn.hidden_size, padding_idx=0) + self.var_node_name_embedding = nn.Embedding( + len(vocab.source), gnn.hidden_size, padding_idx=0 + ) self.node_type_embedder = node_type_embedder self.node_content_embedder = node_content_embedder @@ -52,22 +55,18 @@ def __init__(self, + node_type_embedder.embeddings.embedding_dim + node_content_embedder.embeddings.embedding_dim, gnn.hidden_size, - bias=False + bias=False, ) - self.decoder_cell_init = \ - nn.Linear(gnn.hidden_size, decoder_hidden_size) + self.decoder_cell_init = nn.Linear(gnn.hidden_size, decoder_hidden_size) - self.init_with_seq_encoding = config['init_with_seq_encoding'] + self.init_with_seq_encoding = config["init_with_seq_encoding"] if self.init_with_seq_encoding: - self.seq_encoder = SequentialEncoder.build(config['seq_encoder']) - if config['seq_encoder']['source_encoding_size'] \ - != gnn.hidden_size: - self.seq_variable_encoding_to_graph_linear = \ - nn.Linear( - config['seq_encoder']['source_encoding_size'], - gnn.hidden_size - ) + self.seq_encoder = SequentialEncoder.build(config["seq_encoder"]) + if config["seq_encoder"]["source_encoding_size"] != gnn.hidden_size: + self.seq_variable_encoding_to_graph_linear = nn.Linear( + config["seq_encoder"]["source_encoding_size"], gnn.hidden_size + ) else: self.seq_variable_encoding_to_graph_linear = lambda x: x @@ -80,128 +79,120 @@ def device(self): @classmethod def default_params(cls): return { - 'gnn': { - 'hidden_size': 128, - 'layer_timesteps': [8], - 'residual_connections': {'0': [0]} + "gnn": { + "hidden_size": 128, + "layer_timesteps": [8], + "residual_connections": {"0": [0]}, }, - 'connections': { - 'top_down', - 'bottom_up', - 'variable_master_nodes', - 'terminals', - 'master_node', - 'func_root_to_arg' + "connections": { + "top_down", + "bottom_up", + "variable_master_nodes", + "terminals", + "master_node", + "func_root_to_arg", }, - 'vocab_file': None, - 'bpe_model_path': None, - 'node_syntax_type_embedding_size': 64, - 'node_type_embedding_size': 64, - 'node_content_embedding_size': 128, - 'init_with_seq_encoding': False + "vocab_file": None, + "bpe_model_path": None, + "node_syntax_type_embedding_size": 64, + "node_type_embedding_size": 64, + "node_content_embedding_size": 128, + "init_with_seq_encoding": False, } @classmethod - def build(cls, config): + def build(cls, config, logging=False): params = util.update(GraphASTEncoder.default_params(), config) - if False: + if logging: print(params) - connections = params['connections'] + connections = params["connections"] connection2edge_type = { - 'top_down': 1, - 'bottom_up': 1, - 'variable_master_nodes': 2, - 'terminals': 2, - 'master_node': 2, - 'var_usage': 2, - 'func_root_to_arg': 1 + "top_down": 1, + "bottom_up": 1, + "variable_master_nodes": 2, + "terminals": 2, + "master_node": 2, + "var_usage": 2, + "func_root_to_arg": 1, } num_edge_types = sum(connection2edge_type[key] for key in connections) gnn = GatedGraphNeuralNetwork( - hidden_size=params['gnn']['hidden_size'], - layer_timesteps=params['gnn']['layer_timesteps'], - residual_connections=params['gnn']['residual_connections'], - num_edge_types=num_edge_types + hidden_size=params["gnn"]["hidden_size"], + layer_timesteps=params["gnn"]["layer_timesteps"], + residual_connections=params["gnn"]["residual_connections"], + num_edge_types=num_edge_types, ) - vocab = Vocab.load(params['vocab_file']) + vocab = Vocab.load(params["vocab_file"]) node_type_embedder = NodeTypeEmbedder( - len(vocab.grammar.variable_types), - params['node_type_embedding_size'] + len(vocab.grammar.variable_types), params["node_type_embedding_size"] ) node_content_embedder = SubTokenEmbedder( - vocab.obj_name.subtoken_model_path, - params['node_content_embedding_size'] + vocab.obj_name.subtoken_model_path, params["node_content_embedding_size"] ) - model = cls(gnn, - params['connections'], - params['node_syntax_type_embedding_size'], - params['decoder_hidden_size'], - node_type_embedder, - node_content_embedder, - vocab, - config=params) + model = cls( + gnn, + params["connections"], + params["node_syntax_type_embedding_size"], + params["decoder_hidden_size"], + node_type_embedder, + node_content_embedder, + vocab, + config=params, + ) return model def forward(self, tensor_dict: Dict[str, torch.Tensor]): - tree_node_init_encoding = \ - self.get_batched_tree_init_encoding(tensor_dict) + tree_node_init_encoding = self.get_batched_tree_init_encoding(tensor_dict) if self.init_with_seq_encoding: # scatter sequential encoding results to variable positions - seq_encoding = self.seq_encoder(tensor_dict['seq_encoder_input']) - seq_var_encoding = seq_encoding['variable_encoding'] + seq_encoding = self.seq_encoder(tensor_dict["seq_encoder_input"]) + seq_var_encoding = seq_encoding["variable_encoding"] tree_node_init_encoding[ - tensor_dict['variable_master_node_ids'] + tensor_dict["variable_master_node_ids"] ] = self.seq_variable_encoding_to_graph_linear( - seq_var_encoding.view( - -1, - seq_var_encoding.size(-1) - )[tensor_dict['var_repr_flattened_positions']] + seq_var_encoding.view(-1, seq_var_encoding.size(-1))[ + tensor_dict["var_repr_flattened_positions"] + ] ) # (batch_size, node_encoding_size) tree_node_encoding = self.gnn.compute_node_representations( initial_node_representation=tree_node_init_encoding, - adjacency_lists=tensor_dict['adj_lists'] + adjacency_lists=tensor_dict["adj_lists"], ) - connections = self.config['connections'] + connections = self.config["connections"] - if 'variable_master_nodes' in connections: - variable_master_node_ids = tensor_dict['variable_master_node_ids'] + if "variable_master_nodes" in connections: + variable_master_node_ids = tensor_dict["variable_master_node_ids"] variable_encoding = tree_node_encoding[variable_master_node_ids] else: - var_nodes_encoding = \ - tree_node_encoding[tensor_dict['variable_node_ids']] - variable_num = tensor_dict['variable_mention_nums'].size(0) + var_nodes_encoding = tree_node_encoding[tensor_dict["variable_node_ids"]] + variable_num = tensor_dict["variable_mention_nums"].size(0) variable_encoding = torch.zeros( - variable_num, - var_nodes_encoding.size(-1), - device=self.device + variable_num, var_nodes_encoding.size(-1), device=self.device ) var_nodes_encoding_sum = variable_encoding.scatter_add_( 0, - tensor_dict[ - 'variable_node_variable_ids' - ].unsqueeze(-1).expand_as(var_nodes_encoding), - var_nodes_encoding + tensor_dict["variable_node_variable_ids"] + .unsqueeze(-1) + .expand_as(var_nodes_encoding), + var_nodes_encoding, ) - variable_encoding = \ - var_nodes_encoding_sum \ - / tensor_dict['variable_mention_nums'].unsqueeze(-1) + variable_encoding = var_nodes_encoding_sum / tensor_dict[ + "variable_mention_nums" + ].unsqueeze(-1) # restore variable encoding - variable_encoding = \ - variable_encoding[ - tensor_dict['variable_encoding_restoration_indices'] - ] * tensor_dict[ - 'variable_encoding_restoration_indices_mask' - ].unsqueeze(-1) + variable_encoding = variable_encoding[ + tensor_dict["variable_encoding_restoration_indices"] + ] * tensor_dict["variable_encoding_restoration_indices_mask"].unsqueeze(-1) context_encoding = dict( tree_node_init_encoding=tree_node_init_encoding, @@ -215,11 +206,12 @@ def forward(self, tensor_dict: Dict[str, torch.Tensor]): # noinspection PyUnboundLocalVariable @classmethod - def to_packed_graph(cls, - asts: List[AbstractSyntaxTree], - connections: List, - init_with_seq_encoding: bool = False): - # type: (...) -> Tuple[PackedGraph, Dict] + def to_packed_graph( + cls, + asts: List[AbstractSyntaxTree], + connections: List, + init_with_seq_encoding: bool = False, + ) -> Tuple[PackedGraph, Dict]: packed_graph = PackedGraph(asts) max_variable_num = max(len(ast.variables) for ast in asts) @@ -238,135 +230,133 @@ def to_packed_graph(cls, # list of number of mentions for each variable var_mention_nums = [] - use_variable_master_node = 'variable_master_nodes' in connections - variable_repr_restoration_indices = \ - torch.zeros(len(asts), max_variable_num, dtype=torch.long) - variable_repr_restoration_indices_mask = \ - torch.zeros(len(asts), max_variable_num) + use_variable_master_node = "variable_master_nodes" in connections + variable_repr_restoration_indices = torch.zeros( + len(asts), max_variable_num, dtype=torch.long + ) + variable_repr_restoration_indices_mask = torch.zeros( + len(asts), max_variable_num + ) variable_cum_count = 0 for ast_id, ast in enumerate(asts): for prev_node, succ_node in ast.adjacency_list: - prev_node_packed_id = \ - packed_graph.get_packed_node_id(ast_id, prev_node) - succ_node_packed_id = \ - packed_graph.get_packed_node_id(ast_id, succ_node) + prev_node_packed_id = packed_graph.get_packed_node_id(ast_id, prev_node) + succ_node_packed_id = packed_graph.get_packed_node_id(ast_id, succ_node) - node_adj_list.append(( - prev_node_packed_id, - succ_node_packed_id - )) + node_adj_list.append((prev_node_packed_id, succ_node_packed_id)) - if 'terminals' in connections: + if "terminals" in connections: for i in range(len(ast.terminal_nodes) - 1): terminal_i = ast.terminal_nodes[i] terminal_ip1 = ast.terminal_nodes[i + 1] - terminal_nodes_adj_list.append(( - packed_graph.get_packed_node_id(ast_id, terminal_i), - packed_graph.get_packed_node_id(ast_id, terminal_ip1) - )) + terminal_nodes_adj_list.append( + ( + packed_graph.get_packed_node_id(ast_id, terminal_i), + packed_graph.get_packed_node_id(ast_id, terminal_ip1), + ) + ) # add master node - if 'master_node' in connections: + if "master_node" in connections: master_node_id = packed_graph.register_node( - ast_id, - 'master_node', - group='master_nodes' + ast_id, "master_node", group="master_nodes" ) - master_node_adj_list.extend([ - ( - master_node_id, - packed_graph.get_packed_node_id(ast_id, node) - ) - for node in ast - ]) + master_node_adj_list.extend( + [ + (master_node_id, packed_graph.get_packed_node_id(ast_id, node)) + for node in ast + ] + ) for i, (var_name, var_nodes) in enumerate(ast.variables.items()): var_node_num = len(var_nodes) # add data flow links - if 'var_usage' in connections: + if "var_usage" in connections: for _idx in range(len(var_nodes) - 1): - _var_node, _next_var_node = \ - var_nodes[_idx], var_nodes[_idx + 1] - var_usage_adj_list.append(( - packed_graph.get_packed_node_id(ast_id, - _var_node), - packed_graph.get_packed_node_id(ast_id, - _next_var_node) - )) - - if 'func_root_to_arg' in connections and var_nodes[0].is_arg: - func_root_node_id = \ - packed_graph.get_packed_node_id(ast_id, ast.root) + _var_node, _next_var_node = var_nodes[_idx], var_nodes[_idx + 1] + var_usage_adj_list.append( + ( + packed_graph.get_packed_node_id(ast_id, _var_node), + packed_graph.get_packed_node_id(ast_id, _next_var_node), + ) + ) + + if "func_root_to_arg" in connections and var_nodes[0].is_arg: + func_root_node_id = packed_graph.get_packed_node_id( + ast_id, ast.root + ) for var_node in var_nodes: - func_root_to_arg_adj_list.append(( - func_root_node_id, - packed_graph.get_packed_node_id(ast_id, var_node) - )) + func_root_to_arg_adj_list.append( + ( + func_root_node_id, + packed_graph.get_packed_node_id(ast_id, var_node), + ) + ) if use_variable_master_node: - var_master_node_id, node_id_in_group = \ - packed_graph.register_node( - ast_id, - var_name, - group='variable_master_nodes', - return_node_index_in_group=True - ) + var_master_node_id, node_id_in_group = packed_graph.register_node( + ast_id, + var_name, + group="variable_master_nodes", + return_node_index_in_group=True, + ) - var_master_nodes_adj_list.extend([( - var_master_node_id, - packed_graph.get_packed_node_id(ast_id, node)) - for node in var_nodes - ]) + var_master_nodes_adj_list.extend( + [ + ( + var_master_node_id, + packed_graph.get_packed_node_id(ast_id, node), + ) + for node in var_nodes + ] + ) else: var_node_packed_ids = [ packed_graph.get_packed_node_id(ast_id, node) for node in var_nodes ] var_node_ids.extend(var_node_packed_ids) - var_node_variable_ids.extend( - [variable_cum_count] * var_node_num - ) + var_node_variable_ids.extend([variable_cum_count] * var_node_num) var_mention_nums.append(var_node_num) - variable_repr_restoration_indices[ast_id, i] = \ - variable_cum_count - var_repr_flattened_positions.append( - i + max_variable_num * ast_id - ) + variable_repr_restoration_indices[ast_id, i] = variable_cum_count + var_repr_flattened_positions.append(i + max_variable_num * ast_id) + variable_cum_count += 1 - variable_repr_restoration_indices_mask[ - ast_id, :len(ast.variables) - ] = 1. + + variable_repr_restoration_indices_mask[ast_id, : len(ast.variables)] = 1.0 adj_lists = [] - if 'top_down' in connections: + if "top_down" in connections: adj_lists.append(node_adj_list) - if 'bottom_up' in connections: + if "bottom_up" in connections: reversed_node_adj_list = [(n2, n1) for n1, n2 in node_adj_list] adj_lists.append(reversed_node_adj_list) - if 'variable_master_nodes' in connections: - reversed_var_master_nodes_adj_list = \ - [(n2, n1) for n1, n2 in var_master_nodes_adj_list] + if "variable_master_nodes" in connections: + reversed_var_master_nodes_adj_list = [ + (n2, n1) for n1, n2 in var_master_nodes_adj_list + ] adj_lists.append(master_node_adj_list) adj_lists.append(reversed_var_master_nodes_adj_list) - if 'var_usage' in connections: - reversed_var_usage_adj_list = \ - [(n2, n1) for n1, n2 in var_usage_adj_list] + if "var_usage" in connections: + reversed_var_usage_adj_list = [(n2, n1) for n1, n2 in var_usage_adj_list] adj_lists.append(var_usage_adj_list) adj_lists.append(reversed_var_usage_adj_list) - if 'func_root_to_arg' in connections: + if "func_root_to_arg" in connections: adj_lists.append(func_root_to_arg_adj_list) - if 'terminals' in connections: - reversed_terminal_nodes_adj_list = \ - [(n2, n1) for n1, n2 in terminal_nodes_adj_list] + if "terminals" in connections: + reversed_terminal_nodes_adj_list = [ + (n2, n1) for n1, n2 in terminal_nodes_adj_list + ] adj_lists.append(terminal_nodes_adj_list) adj_lists.append(reversed_terminal_nodes_adj_list) - if 'master_node' in connections: - reversed_master_node_adj_list = \ - [(n2, n1) for n1, n2 in master_node_adj_list] + if "master_node" in connections: + reversed_master_node_adj_list = [ + (n2, n1) for n1, n2 in master_node_adj_list + ] adj_lists.append(master_node_adj_list) adj_lists.append(reversed_master_node_adj_list) @@ -377,35 +367,30 @@ def to_packed_graph(cls, max_tree_node_num = max(tree.size for tree in packed_graph.trees) tree_restoration_indices = torch.zeros( - packed_graph.tree_num, - max_tree_node_num, - dtype=torch.long + packed_graph.tree_num, max_tree_node_num, dtype=torch.long ) tree_restoration_indices_mask = torch.zeros( - packed_graph.tree_num, - max_tree_node_num, - dtype=torch.float + packed_graph.tree_num, max_tree_node_num, dtype=torch.float ) - max_terminal_node_num = max(len(tree.terminal_nodes) - for tree in packed_graph.trees) + max_terminal_node_num = max( + len(tree.terminal_nodes) for tree in packed_graph.trees + ) terminal_node_restoration_indices = torch.zeros( - packed_graph.tree_num, - max_terminal_node_num, - dtype=torch.long + packed_graph.tree_num, max_terminal_node_num, dtype=torch.long ) terminal_node_restoration_indices_mask = torch.zeros( - packed_graph.tree_num, - max_terminal_node_num, - dtype=torch.float + packed_graph.tree_num, max_terminal_node_num, dtype=torch.float ) for ast_id in range(packed_graph.tree_num): - packed_node_ids = \ - list(packed_graph.node_groups[ast_id]['ast_nodes'].values()) - tree_restoration_indices[ast_id, :len(packed_node_ids)] = \ - torch.tensor(packed_node_ids) - tree_restoration_indices_mask[ast_id, :len(packed_node_ids)] = 1. + packed_node_ids = list( + packed_graph.node_groups[ast_id]["ast_nodes"].values() + ) + tree_restoration_indices[ast_id, : len(packed_node_ids)] = torch.tensor( + packed_node_ids + ) + tree_restoration_indices_mask[ast_id, : len(packed_node_ids)] = 1.0 tree = packed_graph.trees[ast_id] terminal_node_ids = [ @@ -413,48 +398,50 @@ def to_packed_graph(cls, for node in tree.terminal_nodes ] terminal_node_restoration_indices[ - ast_id, :len(terminal_node_ids) + ast_id, : len(terminal_node_ids) ] = torch.tensor(terminal_node_ids) terminal_node_restoration_indices_mask[ - ast_id, :len(terminal_node_ids) - ] = 1. + ast_id, : len(terminal_node_ids) + ] = 1.0 tensor_dict = { - 'adj_lists': adj_lists, - 'variable_encoding_restoration_indices': - variable_repr_restoration_indices, - 'variable_encoding_restoration_indices_mask': - variable_repr_restoration_indices_mask, - 'tree_restoration_indices': tree_restoration_indices, - 'tree_restoration_indices_mask': tree_restoration_indices_mask, - 'terminal_node_restoration_indices': - terminal_node_restoration_indices, - 'terminal_node_restoration_indices_mask': - terminal_node_restoration_indices_mask, - 'packed_graph_size': packed_graph.size + "adj_lists": adj_lists, + "variable_encoding_restoration_indices": variable_repr_restoration_indices, + "variable_encoding_restoration_indices_mask": variable_repr_restoration_indices_mask, + "tree_restoration_indices": tree_restoration_indices, + "tree_restoration_indices_mask": tree_restoration_indices_mask, + "terminal_node_restoration_indices": terminal_node_restoration_indices, + "terminal_node_restoration_indices_mask": terminal_node_restoration_indices_mask, + "packed_graph_size": packed_graph.size, } if init_with_seq_encoding: - tensor_dict['var_repr_flattened_positions'] = \ - torch.tensor(var_repr_flattened_positions, dtype=torch.long) + tensor_dict["var_repr_flattened_positions"] = torch.tensor( + var_repr_flattened_positions, dtype=torch.long + ) if use_variable_master_node: - tensor_dict['variable_master_node_ids'] = \ - [_id for node, _id in - packed_graph.get_nodes_by_group('variable_master_nodes')] + tensor_dict["variable_master_node_ids"] = [ + _id + for node, _id in packed_graph.get_nodes_by_group( + "variable_master_nodes" + ) + ] else: - tensor_dict['variable_node_ids'] = torch.tensor(var_node_ids) - tensor_dict['variable_node_variable_ids'] = \ - torch.tensor(var_node_variable_ids) - tensor_dict['variable_mention_nums'] = \ - torch.tensor(var_mention_nums, dtype=torch.float) + tensor_dict["variable_node_ids"] = torch.tensor(var_node_ids) + tensor_dict["variable_node_variable_ids"] = torch.tensor( + var_node_variable_ids + ) + tensor_dict["variable_mention_nums"] = torch.tensor( + var_mention_nums, dtype=torch.float + ) return packed_graph, tensor_dict @classmethod - def to_tensor_dict(cls, packed_graph: PackedGraph, - grammar: Grammar, - vocab: Vocab) -> Dict[str, torch.Tensor]: + def to_tensor_dict( + cls, packed_graph: PackedGraph, grammar: Grammar, vocab: Vocab + ) -> Dict[str, torch.Tensor]: obj_name_bpe_model = vocab.obj_name.subtoken_model obj_name_bpe_pad_idx = vocab.obj_name.subtoken_model.pad_id() @@ -462,23 +449,19 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, variable_master_node_type_idx = len(grammar.syntax_types) master_node_type_idx = variable_master_node_type_idx + 1 - node_syntax_type_indices = \ - torch.zeros(packed_graph.size, dtype=torch.long) - var_node_name_indices = \ - torch.zeros(packed_graph.size, dtype=torch.long) - tree_node_to_tree_id_map = \ - torch.zeros(packed_graph.size, dtype=torch.long) + node_syntax_type_indices = torch.zeros(packed_graph.size, dtype=torch.long) + var_node_name_indices = torch.zeros(packed_graph.size, dtype=torch.long) + tree_node_to_tree_id_map = torch.zeros(packed_graph.size, dtype=torch.long) sub_tokens_list = [] node_with_subtokens_indices = [] node_type_tokens_list = [0 for _ in range(packed_graph.size)] - for i, (ast_id, group, node_key) in \ - enumerate(packed_graph.nodes.values()): + for i, (ast_id, group, node_key) in enumerate(packed_graph.nodes.values()): ast = packed_graph.trees[ast_id] node_type_tokens = [] - if group == 'ast_nodes': + if group == "ast_nodes": node = ast.id_to_node[node_key] type_idx = grammar.syntax_type_to_id[node.node_type] if node.is_variable_node: @@ -486,25 +469,25 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, # function root with type `block` also has an name entry # storing the name of the function - if node.node_type == 'obj' \ - or node.node_type == 'block' \ - and hasattr(node, 'name'): + if ( + node.node_type == "obj" + or node.node_type == "block" + and hasattr(node, "name") + ): # compute variable embedding - node_sub_tokens = \ - obj_name_bpe_model.encode_as_ids(node.name) + node_sub_tokens = obj_name_bpe_model.encode_as_ids(node.name) sub_tokens_list.append(node_sub_tokens) node_with_subtokens_indices.append(i) - if hasattr(node, 'type_tokens'): + if hasattr(node, "type_tokens"): node_type_tokens = [ - vocab.grammar.variable_type_to_id(t) - for t in node.type_tokens + vocab.grammar.variable_type_to_id(t) for t in node.type_tokens ] - elif group == 'variable_master_nodes': + elif group == "variable_master_nodes": type_idx = variable_master_node_type_idx var_node_name_indices[i] = vocab.source[node_key] - elif group == 'master_nodes': + elif group == "master_nodes": type_idx = master_node_type_idx else: raise ValueError() @@ -517,22 +500,20 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, if node_with_subtokens_indices: max_subtoken_num = max(len(x) for x in sub_tokens_list) sub_tokens_indices = np.zeros( - (len(sub_tokens_list), max_subtoken_num), - dtype=np.int64 + (len(sub_tokens_list), max_subtoken_num), dtype=np.int64 ) sub_tokens_indices.fill(obj_name_bpe_pad_idx) for i, token_ids in enumerate(sub_tokens_list): - sub_tokens_indices[i, :len(token_ids)] = token_ids + sub_tokens_indices[i, : len(token_ids)] = token_ids sub_tokens_indices = torch.from_numpy(sub_tokens_indices) max_typetoken_num = max(len(x) for x in node_type_tokens_list) node_type_indices = np.zeros( - (len(node_type_tokens_list), max_typetoken_num), - dtype=np.int64 + (len(node_type_tokens_list), max_typetoken_num), dtype=np.int64 ) for i, type_ids in enumerate(node_type_tokens_list): - node_type_indices[i, :len(type_ids)] = type_ids + node_type_indices[i, : len(type_ids)] = type_ids node_type_indices = torch.from_numpy(node_type_indices) return dict( @@ -542,36 +523,33 @@ def to_tensor_dict(cls, packed_graph: PackedGraph, node_type_indices=node_type_indices, var_node_name_indices=var_node_name_indices, node_with_subtokens_indices=torch.tensor( - node_with_subtokens_indices, - dtype=torch.long + node_with_subtokens_indices, dtype=torch.long ), sub_tokens_indices=sub_tokens_indices, - tree_node_to_tree_id_map=tree_node_to_tree_id_map + tree_node_to_tree_id_map=tree_node_to_tree_id_map, ) - def get_batched_tree_init_encoding(self, - tensor_dict: Dict[str, torch.Tensor]): - # type: (...) -> torch.Tensor + def get_batched_tree_init_encoding( + self, tensor_dict: Dict[str, torch.Tensor] + ) -> torch.Tensor: node_syntax_type_embedding = self.node_syntax_type_embedding( - tensor_dict['node_syntax_type_indices'] + tensor_dict["node_syntax_type_indices"] ) - node_type_embedding = \ - self.node_type_embedder(tensor_dict['node_type_indices']) + node_type_embedding = self.node_type_embedder(tensor_dict["node_type_indices"]) node_content_embedding = self.var_node_name_embedding( - tensor_dict['var_node_name_indices']) * \ - torch.ne( - tensor_dict['var_node_name_indices'], 0. - ).float().unsqueeze(-1) - - if tensor_dict['node_with_subtokens_indices'].size(0) > 0: - obj_node_content_embedding = \ - self.node_content_embedder(tensor_dict['sub_tokens_indices']) + tensor_dict["var_node_name_indices"] + ) * torch.ne(tensor_dict["var_node_name_indices"], 0.0).float().unsqueeze(-1) + + if tensor_dict["node_with_subtokens_indices"].size(0) > 0: + obj_node_content_embedding = self.node_content_embedder( + tensor_dict["sub_tokens_indices"] + ) node_content_embedding = node_content_embedding.scatter( 0, - tensor_dict[ - 'node_with_subtokens_indices' - ].unsqueeze(-1).expand_as(obj_node_content_embedding), - obj_node_content_embedding + tensor_dict["node_with_subtokens_indices"] + .unsqueeze(-1) + .expand_as(obj_node_content_embedding), + obj_node_content_embedding, ) tree_node_embedding = self.type_and_content_hybrid( @@ -579,75 +557,81 @@ def get_batched_tree_init_encoding(self, [ node_syntax_type_embedding, node_type_embedding, - node_content_embedding + node_content_embedding, ], - dim=-1 + dim=-1, ) ) return tree_node_embedding - def unpack_encoding(self, - flattened_node_encodings: torch.Tensor, - batch_syntax_trees: List[AbstractSyntaxTree], - example_node2batch_node_map: Dict): + def unpack_encoding( + self, + flattened_node_encodings: torch.Tensor, + batch_syntax_trees: List[AbstractSyntaxTree], + example_node2batch_node_map: Dict, + ): batch_size = len(batch_syntax_trees) max_node_num = max(tree.size for tree in batch_syntax_trees) index = np.zeros((batch_size, max_node_num), dtype=np.int64) - batch_tree_node_masks = \ - torch.zeros(batch_size, max_node_num, device=self.device) + batch_tree_node_masks = torch.zeros( + batch_size, max_node_num, device=self.device + ) for e_id, syntax_tree in enumerate(batch_syntax_trees): example_nodes_with_batch_id = [ (example_node_id, batch_node_id) - for (_e_id, example_node_id), batch_node_id - in example_node2batch_node_map.items() + for ( + _e_id, + example_node_id, + ), batch_node_id in example_node2batch_node_map.items() if _e_id == e_id ] - sorted_example_nodes_with_batch_id = \ - sorted(example_nodes_with_batch_id, key=lambda t: t[0]) - example_nodes_batch_id = \ - [t[1] for t in sorted_example_nodes_with_batch_id] + sorted_example_nodes_with_batch_id = sorted( + example_nodes_with_batch_id, key=itemgetter(0) + ) + example_nodes_batch_id = [t[1] for t in sorted_example_nodes_with_batch_id] - index[e_id, :len(example_nodes_batch_id)] = example_nodes_batch_id - batch_tree_node_masks[e_id, :len(example_nodes_batch_id)] = 1. + index[e_id, : len(example_nodes_batch_id)] = example_nodes_batch_id + batch_tree_node_masks[e_id, : len(example_nodes_batch_id)] = 1.0 + # (batch_size, max_node_num, node_encoding_size) batch_node_encoding = flattened_node_encodings[ torch.from_numpy(index).to(flattened_node_encodings.device) ] batch_node_encoding.data.masked_fill_( - (1. - batch_tree_node_masks).byte().unsqueeze(-1), 0. + (1.0 - batch_tree_node_masks).byte().unsqueeze(-1), 0.0 ) - return dict(batch_tree_node_encoding=batch_node_encoding, - batch_tree_node_masks=batch_tree_node_masks) + return dict( + batch_tree_node_encoding=batch_node_encoding, + batch_tree_node_masks=batch_tree_node_masks, + ) def get_decoder_init_state(self, context_encoding, config=None): # compute initial decoder's state via average pooling # (packed_graph_size, encoding_size) - packed_tree_node_encoding = \ - context_encoding['packed_tree_node_encoding'] + packed_tree_node_encoding = context_encoding["packed_tree_node_encoding"] - tree_num = context_encoding['tree_num'] - total_node_num = context_encoding['tree_node_to_tree_id_map'].size(0) + tree_num = context_encoding["tree_num"] + total_node_num = context_encoding["tree_node_to_tree_id_map"].size(0) encoding_size = packed_tree_node_encoding.size(-1) - zero_encoding = \ - packed_tree_node_encoding.new_zeros(tree_num, encoding_size) + + zero_encoding = packed_tree_node_encoding.new_zeros(tree_num, encoding_size) node_encoding_sum = zero_encoding.scatter_add_( 0, - context_encoding['tree_node_to_tree_id_map'].unsqueeze(-1).expand( - -1, - encoding_size - ), - packed_tree_node_encoding) - tree_node_num = \ - packed_tree_node_encoding.new_zeros(tree_num).scatter_add_( - 0, - context_encoding['tree_node_to_tree_id_map'], - packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.) - ) + context_encoding["tree_node_to_tree_id_map"] + .unsqueeze(-1) + .expand(-1, encoding_size), + packed_tree_node_encoding, + ) + tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_( + 0, + context_encoding["tree_node_to_tree_id_map"], + packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.0), + ) avg_node_encoding = node_encoding_sum / tree_node_num.unsqueeze(-1) c_0 = self.decoder_cell_init(avg_node_encoding) @@ -655,22 +639,17 @@ def get_decoder_init_state(self, context_encoding, config=None): return h_0, c_0 - def get_attention_memory(self, - context_encoding, - att_target='terminal_nodes'): - packed_tree_node_enc = \ - context_encoding['packed_tree_node_encoding'] - if att_target == 'ast_nodes': - memory = packed_tree_node_enc[ - context_encoding['tree_restoration_indices'] - ] - mask = context_encoding['tree_restoration_indices_mask'] - elif att_target == 'terminal_nodes': + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + packed_tree_node_enc = context_encoding["packed_tree_node_encoding"] + if att_target == "ast_nodes": + memory = packed_tree_node_enc[context_encoding["tree_restoration_indices"]] + mask = context_encoding["tree_restoration_indices_mask"] + elif att_target == "terminal_nodes": memory = packed_tree_node_enc[ - context_encoding['terminal_node_restoration_indices'] + context_encoding["terminal_node_restoration_indices"] ] - mask = context_encoding['terminal_node_restoration_indices_mask'] + mask = context_encoding["terminal_node_restoration_indices_mask"] else: - raise ValueError('unknown attention target') + raise ValueError("unknown attention target") return memory, mask diff --git a/prediction-plugin/model/hybrid_encoder.py b/prediction-plugin/model/hybrid_encoder.py index c3a9735..9ba9d5f 100644 --- a/prediction-plugin/model/hybrid_encoder.py +++ b/prediction-plugin/model/hybrid_encoder.py @@ -1,10 +1,9 @@ from typing import Dict import torch -from torch import nn as nn - from model.graph_encoder import GraphASTEncoder from model.sequential_encoder import SequentialEncoder +from torch import nn as nn from utils import util @@ -12,18 +11,19 @@ class HybridEncoder(nn.Module): def __init__(self, config): super(HybridEncoder, self).__init__() - self.graph_encoder = GraphASTEncoder.build(config['graph_encoder']) - self.seq_encoder = SequentialEncoder.build(config['seq_encoder']) + self.graph_encoder = GraphASTEncoder.build(config["graph_encoder"]) + self.seq_encoder = SequentialEncoder.build(config["seq_encoder"]) - self.hybrid_method = config['hybrid_method'] - if self.hybrid_method == 'linear_proj': + self.hybrid_method = config["hybrid_method"] + if self.hybrid_method == "linear_proj": self.projection = nn.Linear( - config['seq_encoder']['source_encoding_size'] - + config['graph_encoder']['gnn']['hidden_size'], - config['source_encoding_size'], bias=False + config["seq_encoder"]["source_encoding_size"] + + config["graph_encoder"]["gnn"]["hidden_size"], + config["source_encoding_size"], + bias=False, ) else: - assert self.hybrid_method == 'concat' + assert self.hybrid_method == "concat" self.config = config @@ -36,7 +36,7 @@ def default_params(cls): return { "graph_encoder": GraphASTEncoder.default_params(), "seq_encoder": SequentialEncoder.default_params(), - "hybrid_method": "linear_proj" + "hybrid_method": "linear_proj", } @classmethod @@ -46,58 +46,55 @@ def build(cls, config): return cls(params) def forward(self, tensor_dict: Dict[str, torch.Tensor]): - graph_encoding = self.graph_encoder(tensor_dict['graph_encoder_input']) - seq_encoding = self.seq_encoder(tensor_dict['seq_encoder_input']) + graph_encoding = self.graph_encoder(tensor_dict["graph_encoder_input"]) + seq_encoding = self.seq_encoder(tensor_dict["seq_encoder_input"]) - graph_var_encoding = graph_encoding['variable_encoding'] - seq_var_encoding = seq_encoding['variable_encoding'] + graph_var_encoding = graph_encoding["variable_encoding"] + seq_var_encoding = seq_encoding["variable_encoding"] - if self.hybrid_method == 'linear_proj': + if self.hybrid_method == "linear_proj": variable_encoding = self.projection( - torch.cat([graph_var_encoding, seq_var_encoding], dim=-1)) - else: - variable_encoding = \ torch.cat([graph_var_encoding, seq_var_encoding], dim=-1) + ) + else: + variable_encoding = torch.cat( + [graph_var_encoding, seq_var_encoding], dim=-1 + ) context_encoding = dict( - batch_size=tensor_dict['batch_size'], + batch_size=tensor_dict["batch_size"], variable_encoding=variable_encoding, graph_encoding_result=graph_encoding, - seq_encoding_result=seq_encoding + seq_encoding_result=seq_encoding, ) return context_encoding def get_decoder_init_state(self, context_encoding, config=None): - gnn_dec_init_state, gnn_dec_init_cell = \ - self.graph_encoder.get_decoder_init_state( - context_encoding['graph_encoding_result'] - ) - seq_dec_init_state, seq_dec_init_cell = \ - self.seq_encoder.get_decoder_init_state( - context_encoding['seq_encoding_result'] - ) + ( + gnn_dec_init_state, + gnn_dec_init_cell, + ) = self.graph_encoder.get_decoder_init_state( + context_encoding["graph_encoding_result"] + ) + seq_dec_init_state, seq_dec_init_cell = self.seq_encoder.get_decoder_init_state( + context_encoding["seq_encoding_result"] + ) dec_init_state = (gnn_dec_init_state + seq_dec_init_state) / 2 dec_init_cell = (gnn_dec_init_cell + seq_dec_init_cell) / 2 return dec_init_state, dec_init_cell - def get_attention_memory(self, - context_encoding, - att_target='terminal_nodes'): - assert att_target == 'terminal_nodes' + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + assert att_target == "terminal_nodes" - seq_memory, seq_mask = \ - self.seq_encoder.get_attention_memory( - context_encoding['seq_encoding_result'], - att_target='terminal_nodes' - ) - gnn_memory, gnn_mask = \ - self.graph_encoder.get_attention_memory( - context_encoding['graph_encoding_result'], - att_target='terminal_nodes' - ) + seq_memory, seq_mask = self.seq_encoder.get_attention_memory( + context_encoding["seq_encoding_result"], att_target="terminal_nodes" + ) + gnn_memory, gnn_mask = self.graph_encoder.get_attention_memory( + context_encoding["graph_encoding_result"], att_target="terminal_nodes" + ) if gnn_memory.size(-1) < seq_memory.size(-1): # pad values @@ -106,11 +103,12 @@ def get_attention_memory(self, gnn_memory.new_zeros( gnn_memory.size(0), gnn_memory.size(1), - seq_memory.size(-1) - gnn_memory.size(-1) + seq_memory.size(-1) - gnn_memory.size(-1), ), - gnn_memory + gnn_memory, ], - dim=-1) + dim=-1, + ) memory = torch.cat([gnn_memory, seq_memory], dim=1) mask = torch.cat([gnn_mask, seq_mask], dim=1) diff --git a/prediction-plugin/model/model.py b/prediction-plugin/model/model.py index 2c51c8f..1d95fb0 100644 --- a/prediction-plugin/model/model.py +++ b/prediction-plugin/model/model.py @@ -1,63 +1,59 @@ import pprint import sys -from typing import List, Dict, Tuple, Iterable +from typing import Dict, Iterable, List, Tuple import torch import torch.nn as nn - -from utils import util from model.decoder import Decoder -from model.attentional_recurrent_subtoken_decoder \ - import AttentionalRecurrentSubtokenDecoder from model.encoder import Encoder +from model.graph_encoder import GraphASTEncoder from model.hybrid_encoder import HybridEncoder +from model.sequential_encoder import SequentialEncoder +from utils import util from utils.dataset import Batcher, Example class RenamingModel(nn.Module): def __init__(self, encoder: Encoder, decoder: Decoder): - super(RenamingModel, self).__init__() + super().__init__() + self.encoder = encoder self.decoder = decoder self.config: Dict = None + @property + def vocab(self): + return self.encoder.vocab + @property def batcher(self): - if not hasattr(self, '_batcher'): + if not hasattr(self, "_batcher"): _batcher = Batcher(self.config) - setattr(self, '_batcher', _batcher) + setattr(self, "_batcher", _batcher) + return self._batcher @property def device(self): return self.encoder.device - @property - def vocab(self): - return self.encoder.vocab - @staticmethod def default_params(): return { - 'train': { - 'unchanged_variable_weight': 1.0, - 'max_epoch': 30, - 'patience': 5 - }, - 'decoder': { - 'type': 'SimpleDecoder' - } + "train": {"unchanged_variable_weight": 1.0, "max_epoch": 30, "patience": 5}, + "decoder": {"type": "SimpleDecoder"}, } @classmethod - def build(cls, config): - params = util.update(RenamingModel.default_params(), config) - encoder = globals()[config['encoder']['type']].build(config['encoder']) - decoder = globals()[config['decoder']['type']].build(config['decoder']) + def build(cls, config, logging=False): + params = util.update(cls.default_params(), config) + encoder = globals()[config["encoder"]["type"]].build(config["encoder"]) + decoder = globals()[config["decoder"]["type"]].build(config["decoder"]) model = cls(encoder, decoder) - params = util.update(params, {'encoder': encoder.config, - 'decoder': decoder.config}) + params = util.update( + params, {"encoder": encoder.config, "decoder": decoder.config} + ) model.config = params # give the decoder a reference to the encoder model.decoder.encoder = encoder @@ -66,18 +62,20 @@ def build(cls, config): encoder.batcher = model.batcher decoder.batcher = model.batcher - if False: - print('Current Configuration:', file=sys.stderr) + if logging: + print("Current Configuration:", file=sys.stderr) pp = pprint.PrettyPrinter(indent=2, stream=sys.stderr) pp.pprint(model.config) return model - def forward(self, - source_asts: Dict[str, torch.Tensor], - prediction_target: Dict[str, torch.Tensor]): - # type: (...) -> Tuple[torch.Tensor, Dict] - """Given a batch of decompiled abstract syntax trees, and the + def forward( + self, + source_asts: Dict[str, torch.Tensor], + prediction_target: Dict[str, torch.Tensor], + ) -> Tuple[torch.Tensor, Dict]: + """ + Given a batch of decompiled abstract syntax trees, and the gold-standard renaming of variable nodes, compute the log-likelihood of the gold-standard renaming for training @@ -88,36 +86,45 @@ def forward(self, Return: tensor of size batch_size denoting the log-likelihood of renamings - """ + + # src_ast_encoding: (batch_size, max_ast_node_num, node_encoding_size) + # src_ast_mask: (batch_size, max_ast_node_num) context_encoding = self.encoder(source_asts) + + # (batch_size, variable_num, vocab_size) or (prediction_node_num, vocab_size) var_name_log_probs = self.decoder(context_encoding, prediction_target) - result = self.decoder.get_target_log_prob(var_name_log_probs, - prediction_target, - context_encoding) - tgt_var_name_log_prob = result['tgt_var_name_log_prob'] - tgt_weight = prediction_target['variable_tgt_name_weight'] + + result = self.decoder.get_target_log_prob( + var_name_log_probs, prediction_target, context_encoding + ) + tgt_var_name_log_prob = result["tgt_var_name_log_prob"] + tgt_weight = prediction_target["variable_tgt_name_weight"] weighted_log_prob = tgt_var_name_log_prob * tgt_weight - tgt_var_encoding_indices_mask = \ - prediction_target['target_variable_encoding_indices_mask'] - ast_log_probs = \ - weighted_log_prob.sum(dim=-1) \ - / tgt_var_encoding_indices_mask.sum(-1) - result['batch_log_prob'] = ast_log_probs + tgt_var_encoding_indices_mask = prediction_target[ + "target_variable_encoding_indices_mask" + ] + ast_log_probs = weighted_log_prob.sum( + dim=-1 + ) / tgt_var_encoding_indices_mask.sum(-1) + result["batch_log_prob"] = ast_log_probs + return result - def decode_dataset(self, dataset, batch_size=4096): - # type: (...) -> Iterable[Tuple[Example, Dict]] + def decode_dataset( + self, dataset, batch_size=4096 + ) -> Iterable[Tuple[Example, Dict]]: with torch.no_grad(): - data_iter = dataset.batch_iterator(batch_size=batch_size, - train=False, - progress=False, - config=self.config) + data_iter = dataset.batch_iterator( + batch_size=batch_size, train=False, progress=False, config=self.config + ) self.eval() + for batch in data_iter: - rename_results = \ - self.decoder.predict([e.ast for e in batch.examples], - self.encoder) + rename_results = self.decoder.predict( + [e.ast for e in batch.examples], self.encoder + ) + for example, result in zip(batch.examples, rename_results): yield example, result @@ -126,24 +133,27 @@ def predict(self, examples: List[Example]): def save(self, model_path, **kwargs): params = { - 'config': self.config, - 'state_dict': self.state_dict(), - 'kwargs': kwargs + "config": self.config, + "state_dict": self.state_dict(), + "kwargs": kwargs, } + torch.save(params, model_path) @classmethod - def load(cls, model_path, use_cuda=False, new_config=None): - # type: (...) -> 'RenamingModel' + def load(cls, model_path, use_cuda=False, new_config=None) -> "RenamingModel": device = torch.device("cuda:0" if use_cuda else "cpu") - params = torch.load(model_path, map_location=lambda store, _: store) - if new_config is not None: - config = util.update(params['config'], new_config) - kwargs = dict() - if params['kwargs'] is not None: - kwargs = params['kwargs'] + params = torch.load(model_path, map_location=lambda storage, _loc: storage) + print('\n\n\n') + pprint.pprint(params["config"]) + print('\n\n\n') + + config = util.update(params["config"], new_config) + + kwargs = dict() if params["kwargs"] is None else params["kwargs"] + model = cls.build(config, **kwargs) - model.load_state_dict(params['state_dict'], strict=False) + model.load_state_dict(params["state_dict"], strict=False) model = model.to(device) model.eval() diff --git a/prediction-plugin/model/recurrent_decoder.py b/prediction-plugin/model/recurrent_decoder.py new file mode 100644 index 0000000..a51ce84 --- /dev/null +++ b/prediction-plugin/model/recurrent_decoder.py @@ -0,0 +1,351 @@ +from collections import OrderedDict, namedtuple +from typing import Dict, List + +import torch +from model.decoder import Decoder +from model.encoder import Encoder +from torch import nn as nn +from utils import nn_util, util +from utils.ast import AbstractSyntaxTree +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab + + +class RecurrentDecoder(Decoder): + def __init__( + self, + ast_node_encoding_size: int, + hidden_size: int, + dropout: float, + vocab: Vocab, + ): + super(Decoder, self).__init__() + + self.vocab = vocab + + self.lstm_cell = nn.LSTMCell( + ast_node_encoding_size + ast_node_encoding_size, hidden_size + ) + self.decoder_cell_init = nn.Linear(ast_node_encoding_size, hidden_size) + self.state2names = nn.Linear( + ast_node_encoding_size, len(vocab.target), bias=True + ) + self.dropout = nn.Dropout(dropout) + self.config: Dict = None + + self.Hypothesis = namedtuple("Hypothesis", ["variable_list", "score"]) + + @property + def device(self): + return self.state2names.weight.device + + @classmethod + def default_params(cls): + return { + "vocab_file": None, + "ast_node_encoding_size": 128, + "hidden_size": 128, + "input_feed": False, + "dropout": 0.2, + "beam_size": 5, + "unk_replace": True, + } + + @classmethod + def build(cls, config): + params = util.update(cls.default_params(), config) + + vocab = Vocab.load(params["vocab_file"]) + model = cls( + params["ast_node_encoding_size"], + params["hidden_size"], + params["dropout"], + vocab, + ) + model.config = params + + return model + + def get_init_state(self, src_ast_encoding): + # compute initial decoder's state via average pooling + + # (packed_graph_size, encoding_size) + packed_tree_node_encoding = src_ast_encoding["packed_tree_node_encoding"] + + tree_num = src_ast_encoding["tree_num"] + total_node_num = src_ast_encoding["tree_node_to_tree_id_map"].size(0) + encoding_size = packed_tree_node_encoding.size(-1) + + zero_encoding = packed_tree_node_encoding.new_zeros(tree_num, encoding_size) + + node_encoding_sum = zero_encoding.scatter_add_( + 0, + src_ast_encoding["tree_node_to_tree_id_map"] + .unsqueeze(-1) + .expand(-1, encoding_size), + packed_tree_node_encoding, + ) + tree_node_num = packed_tree_node_encoding.new_zeros(tree_num).scatter_add_( + 0, + src_ast_encoding["tree_node_to_tree_id_map"], + packed_tree_node_encoding.new_zeros(total_node_num).fill_(1.0), + ) + avg_node_encoding = node_encoding_sum / tree_node_num.unsqueeze(-1) + + c_0 = self.decoder_cell_init(avg_node_encoding) + h_0 = torch.tanh(c_0) + + return h_0, c_0 + + def rnn_step(self, x, h_tm1, src_ast_encoding): + h_t = self.lstm_cell(x, h_tm1) + # TODO: implement attention? + # att_t = torch.tanh(self.att_vec_linear(torch.cat([h_t], 1))) + q_t = self.dropout(h_t[0]) + + return h_t, q_t, None + + def forward(self, src_ast_encoding, prediction_target): + # (prediction_node_num, encoding_size) + variable_master_node_encoding = src_ast_encoding[ + "variable_master_node_encoding" + ] + # (batch_size, max_prediction_node_num) + variable_master_node_restoration_indices = src_ast_encoding[ + "variable_encoding_restoration_indices" + ] + variable_master_node_restoration_indices_mask = src_ast_encoding[ + "variable_encoding_restoration_indices_mask" + ] + + # (batch_size, max_prediction_node_num, encoding_size) + variable_master_node_encoding = variable_master_node_encoding[ + variable_master_node_restoration_indices + ] + variable_tgt_name_id = prediction_target["variable_tgt_name_id"][ + variable_master_node_restoration_indices + ] + + batch_size = variable_tgt_name_id.size(0) + variable_encoding_size = variable_master_node_encoding.size(-1) + + h_0 = self.get_init_state(src_ast_encoding) + att_tm1 = variable_master_node_encoding.new_zeros( + src_ast_encoding["tree_num"], variable_master_node_encoding.size(-1) + ) + + h_tm1 = h_0 + query_vecs = [] + for t, variable_encoding_t in enumerate( + variable_master_node_encoding.split(split_size=1, dim=1) + ): + # variable_encoding_t: (batch_size, encoding_size) + variable_encoding_t = variable_encoding_t.squeeze(1) + if t > 0: + # (batch_size) + v_tm1_name_id = variable_tgt_name_id[:, t - 1] + # (batch_size, encoding_size) + v_tm1_name_embed = self.state2names.weight[v_tm1_name_id] + else: + # (batch_size, encoding_size) + v_tm1_name_embed = torch.zeros( + batch_size, variable_encoding_size, device=self.device + ) + + if self.config["input_feed"]: + x = torch.cat([variable_encoding_t, v_tm1_name_embed, att_tm1], dim=-1) + else: + x = torch.cat([variable_encoding_t, v_tm1_name_embed], dim=-1) + + h_t, q_t, alpha_t = self.rnn_step(x, h_tm1, src_ast_encoding) + + att_tm1 = q_t + h_tm1 = h_t + query_vecs.append(q_t) + + # (batch_size, max_prediction_node_num, encoding_size) + query_vecs = torch.stack(query_vecs).permute(1, 0, 2) + + # (batch_size, max_prediction_node_num, vocab_size) + logits = self.state2names(query_vecs) + var_name_log_probs = torch.log_softmax(logits, dim=-1) + var_name_log_probs = ( + var_name_log_probs + * variable_master_node_restoration_indices_mask.unsqueeze(-1) + ) + + return var_name_log_probs + + def get_target_log_prob( + self, var_name_log_probs, prediction_target, src_ast_encoding + ): + # (batch_size, max_prediction_node_num) + variable_tgt_name_id = prediction_target["variable_tgt_name_id"][ + src_ast_encoding["variable_encoding_restoration_indices"] + ] + tgt_var_name_log_prob = torch.gather( + var_name_log_probs, dim=-1, index=variable_tgt_name_id.unsqueeze(-1) + ).squeeze(-1) + + tgt_var_name_log_prob = ( + tgt_var_name_log_prob + * src_ast_encoding["variable_encoding_restoration_indices_mask"] + ) + + result = dict(tgt_var_name_log_prob=tgt_var_name_log_prob) + + return result + + def predict( + self, source_asts: List[AbstractSyntaxTree], encoder: Encoder + ) -> List[Dict]: + beam_size = self.config["beam_size"] + unk_replace = self.config["unk_replace"] + + variable_nums = [] + for ast_id, ast in enumerate(source_asts): + variable_nums.append(len(ast.variables)) + + beams = OrderedDict( + (ast_id, [self.Hypothesis([], 0.0)]) + for ast_id, ast in enumerate(source_asts) + ) + hyp_scores_tm1 = torch.zeros(len(beams), 1, device=self.device) + completed_hyps = [[] for _ in source_asts] + tgt_vocab_size = len(self.vocab.target) + + tensor_dict = self.batcher.to_tensor_dict(source_asts=source_asts) + nn_util.to(tensor_dict, self.device) + + context_encoding = encoder(tensor_dict) + h_tm1 = h_0 = self.get_init_state(context_encoding) + + # (prediction_node_num, encoding_size) + variable_master_node_encoding = context_encoding[ + "variable_master_node_encoding" + ] + encoding_size = variable_master_node_encoding.size(1) + # (batch_size, max_prediction_node_num) + variable_master_node_restoration_indices = context_encoding[ + "variable_encoding_restoration_indices" + ] + + # (batch_size, max_prediction_node_num, encoding_size) + variable_master_node_encoding = variable_master_node_encoding[ + variable_master_node_restoration_indices + ] + variable_encoding_t = variable_master_node_encoding[:, 0] + # (batch_size, encoding_size) + variable_name_embed_tm1 = att_tm1 = torch.zeros( + len(source_asts), encoding_size, device=self.device + ) + + max_prediction_node_num = variable_master_node_encoding.size(1) + + for t in range(0, max_prediction_node_num): + live_beam_size = beam_size if t > 0 else 1 + live_tree_ids = [ast_id for ast_id in beams] + + if t > 0: + variable_encoding_t = ( + variable_master_node_encoding[live_tree_ids][:, t] + .unsqueeze(1) + .expand(-1, beam_size, -1) + .contiguous() + .view(-1, encoding_size) + ) + + if self.config["input_feed"]: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 + ) + else: + x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) + + h_t, q_t, alpha_t = self.rnn_step(x, h_tm1, context_encoding) + + # (live_beam_num, beam_size, encoding_size) + q_t_by_beam = q_t.view(len(beams), -1, q_t.size(-1)) + # (live_beam_num, beam_size, vocab_size) + hyp_var_name_scores_t = torch.log_softmax( + self.state2names(q_t_by_beam), dim=-1 + ) + + if unk_replace: + hyp_var_name_scores_t[:, :, self.vocab.target[""]] = float("-inf") + + cont_cand_hyp_scores = hyp_scores_tm1.unsqueeze(-1) + hyp_var_name_scores_t + + # (live_beam_num, beam_size) + new_hyp_scores, new_hyp_position_list = torch.topk( + cont_cand_hyp_scores.view(len(beams), -1), k=beam_size, dim=-1 + ) + + # (live_beam_num, beam_size) + prev_hyp_ids = new_hyp_position_list / tgt_vocab_size + hyp_var_name_ids = new_hyp_position_list % tgt_vocab_size + new_hyp_scores = new_hyp_scores + + # move this tensor to cpu for fast indexing + _prev_hyp_ids = prev_hyp_ids.cpu() + _hyp_var_name_ids = hyp_var_name_ids.cpu() + _new_hyp_scores = new_hyp_scores.cpu() + + new_beams = OrderedDict() + live_beam_ids = [] + for beam_id, (ast_id, beam) in enumerate(beams.items()): + new_hyps = [] + for i in range(beam_size): + prev_hyp_id = _prev_hyp_ids[beam_id, i].item() + prev_hyp = beam[prev_hyp_id] + hyp_var_name_id = _hyp_var_name_ids[beam_id, i].item() + new_hyp_score = _new_hyp_scores[beam_id, i].item() + + new_hyp = self.Hypothesis( + variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], + score=new_hyp_score, + ) + new_hyps.append(new_hyp) + if t + 1 == variable_nums[ast_id]: + completed_hyps[ast_id] = new_hyps + else: + new_beams[ast_id] = new_hyps + live_beam_ids.append(beam_id) + + if t < max_prediction_node_num - 1: + # (live_beam_num, beam_size, *) + prev_hyp_ids = ( + torch.arange(len(beams)).to(self.device) * live_beam_size + ).unsqueeze(-1) + prev_hyp_ids + live_prev_hyp_ids = prev_hyp_ids[live_beam_ids].view(-1) + live_beam_ids = torch.tensor(live_beam_ids, device=self.device) + + hyp_scores_tm1 = new_hyp_scores[live_beam_ids] + h_tm1 = (h_t[0][live_prev_hyp_ids], h_t[1][live_prev_hyp_ids]) + att_tm1 = q_t[live_prev_hyp_ids] + + # (live_beam_num * beam_size) + live_hyp_var_name_ids = hyp_var_name_ids[live_beam_ids].view(-1) + # (live_beam_num * beam_size, embed_size) + variable_name_embed_tm1 = self.state2names.weight[live_hyp_var_name_ids] + + beams = new_beams + + variable_rename_results = [] + for i, hyps in enumerate(completed_hyps): + ast = source_asts[i] + hyps = sorted(hyps, key=lambda hyp: -hyp.score) + top_hyp = hyps[0] + variable_rename_result = dict() + for old_name, var_name_id in zip(ast.variables, top_hyp.variable_list): + new_var_name = self.vocab.target.id2word[var_name_id] + if new_var_name == SAME_VARIABLE_TOKEN: + new_var_name = old_name + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } + + variable_rename_results.append(variable_rename_result) + + return variable_rename_results diff --git a/prediction-plugin/model/recurrent_subtoken_decoder.py b/prediction-plugin/model/recurrent_subtoken_decoder.py index d000d5a..ebf979b 100644 --- a/prediction-plugin/model/recurrent_subtoken_decoder.py +++ b/prediction-plugin/model/recurrent_subtoken_decoder.py @@ -1,20 +1,27 @@ import sys -from collections import namedtuple, OrderedDict +from collections import OrderedDict, namedtuple from typing import Dict, List import torch -from torch import nn as nn - from model.decoder import Decoder from model.encoder import Encoder -from utils import util, nn_util +from torch import nn as nn +from utils import nn_util, util from utils.ast import AbstractSyntaxTree from utils.dataset import Example -from utils.vocab import Vocab, SAME_VARIABLE_TOKEN, END_OF_VARIABLE_TOKEN +from utils.vocab import END_OF_VARIABLE_TOKEN, SAME_VARIABLE_TOKEN, Vocab class RecurrentSubtokenDecoder(Decoder): - def __init__(self, variable_encoding_size: int, hidden_size: int, dropout: float, tie_embed: bool, input_feed: bool, vocab: Vocab): + def __init__( + self, + variable_encoding_size: int, + hidden_size: int, + dropout: float, + tie_embed: bool, + input_feed: bool, + vocab: Vocab, + ): super(Decoder, self).__init__() self.vocab = vocab @@ -31,7 +38,9 @@ def __init__(self, variable_encoding_size: int, hidden_size: int, dropout: float self.dropout = nn.Dropout(dropout) self.config: Dict = None - self.Hypothesis = namedtuple('Hypothesis', ['variable_list', 'variable_ptr', 'score']) + self.Hypothesis = namedtuple( + "Hypothesis", ["variable_list", "variable_ptr", "score"] + ) @property def device(self): @@ -40,28 +49,34 @@ def device(self): @classmethod def default_params(cls): return { - 'vocab_file': None, - 'variable_encoding_size': 128, - 'hidden_size': 128, - 'input_feed': False, - 'tie_embedding': True, - 'dropout': 0.2, - 'beam_size': 5, - 'max_prediction_time_step': 1200, - 'independent_prediction_for_each_variable': False + "vocab_file": None, + "variable_encoding_size": 128, + "hidden_size": 128, + "input_feed": False, + "tie_embedding": True, + "dropout": 0.2, + "beam_size": 5, + "max_prediction_time_step": 1200, + "independent_prediction_for_each_variable": False, } @property def independent_prediction_for_each_variable(self): - return self.config['independent_prediction_for_each_variable'] + return self.config["independent_prediction_for_each_variable"] @classmethod def build(cls, config): params = util.update(cls.default_params(), config) - vocab = Vocab.load(params['vocab_file']) - model = cls(params['variable_encoding_size'], - params['hidden_size'], params['dropout'], params['tie_embedding'], params['input_feed'], vocab) + vocab = Vocab.load(params["vocab_file"]) + model = cls( + params["variable_encoding_size"], + params["hidden_size"], + params["dropout"], + params["tie_embedding"], + params["input_feed"], + vocab, + ) model.config = params return model @@ -79,31 +94,46 @@ def rnn_step(self, x, h_tm1, src_ast_encoding): def forward(self, src_ast_encoding, prediction_target): # (batch_size, max_time_step) - target_variable_encoding_indices = prediction_target['target_variable_encoding_indices'] - target_variable_encoding_indices_mask = prediction_target['target_variable_encoding_indices_mask'] + target_variable_encoding_indices = prediction_target[ + "target_variable_encoding_indices" + ] + target_variable_encoding_indices_mask = prediction_target[ + "target_variable_encoding_indices_mask" + ] batch_size = target_variable_encoding_indices.size(0) - variable_encoding_size = src_ast_encoding['variable_encoding'].size(-1) + variable_encoding_size = src_ast_encoding["variable_encoding"].size(-1) # (batch_size, max_time_step, encoding_size) # scatter variable encoding to sub-token time steps - variable_encoding = torch.gather(src_ast_encoding['variable_encoding'], 1, - target_variable_encoding_indices.unsqueeze(-1).expand(-1, -1, variable_encoding_size)) + variable_encoding = torch.gather( + src_ast_encoding["variable_encoding"], + 1, + target_variable_encoding_indices.unsqueeze(-1).expand( + -1, -1, variable_encoding_size + ), + ) # (batch_size, max_time_step, encoding_size) - variable_tgt_name_id = prediction_target['variable_tgt_name_id'] + variable_tgt_name_id = prediction_target["variable_tgt_name_id"] h_0 = self.get_init_state(src_ast_encoding) - att_tm1 = variable_encoding.new_zeros(src_ast_encoding['batch_size'], self.lstm_cell.hidden_size) - v_tm1_name_embed = torch.zeros(batch_size, self.lstm_cell.hidden_size, device=self.device) + att_tm1 = variable_encoding.new_zeros( + src_ast_encoding["batch_size"], self.lstm_cell.hidden_size + ) + v_tm1_name_embed = torch.zeros( + batch_size, self.lstm_cell.hidden_size, device=self.device + ) h_tm1 = h_0 query_vecs = [] max_time_step = variable_encoding.size(1) - for t, variable_encoding_t in enumerate(variable_encoding.split(split_size=1, dim=1)): + for t, variable_encoding_t in enumerate( + variable_encoding.split(split_size=1, dim=1) + ): # variable_encoding_t: (batch_size, encoding_size) variable_encoding_t = variable_encoding_t.squeeze(1) - if self.config['input_feed']: + if self.config["input_feed"]: x = torch.cat([variable_encoding_t, v_tm1_name_embed, att_tm1], dim=-1) else: x = torch.cat([variable_encoding_t, v_tm1_name_embed], dim=-1) @@ -114,7 +144,7 @@ def forward(self, src_ast_encoding, prediction_target): h_tm1 = h_t query_vecs.append(q_t) v_tm1_name_id = variable_tgt_name_id[:, t] - if self.config['tie_embedding']: + if self.config["tie_embedding"]: v_tm1_name_embed = self.state2names.weight[v_tm1_name_id] else: v_tm1_name_embed = self.var_name_embed(v_tm1_name_id) @@ -124,8 +154,13 @@ def forward(self, src_ast_encoding, prediction_target): variable_ids_tp1 = target_variable_encoding_indices[:, t + 1] variable_ids_t = target_variable_encoding_indices[:, t] - is_tp1_same_variable = torch.eq(variable_ids_tp1, variable_ids_t).float().unsqueeze(-1) # TODO: check if correct! - h_tm1 = (h_tm1[0] * is_tp1_same_variable, h_tm1[1] * is_tp1_same_variable) + is_tp1_same_variable = ( + torch.eq(variable_ids_tp1, variable_ids_t).float().unsqueeze(-1) + ) # TODO: check if correct! + h_tm1 = ( + h_tm1[0] * is_tp1_same_variable, + h_tm1[1] * is_tp1_same_variable, + ) att_tm1 = att_tm1 * is_tp1_same_variable v_tm1_name_embed = v_tm1_name_embed * is_tp1_same_variable @@ -135,17 +170,25 @@ def forward(self, src_ast_encoding, prediction_target): # (batch_size, max_prediction_node_num, vocab_size) logits = self.state2names(query_vecs) var_name_log_probs = torch.log_softmax(logits, dim=-1) - var_name_log_probs = var_name_log_probs * target_variable_encoding_indices_mask.unsqueeze(-1) + var_name_log_probs = ( + var_name_log_probs * target_variable_encoding_indices_mask.unsqueeze(-1) + ) return var_name_log_probs - def get_target_log_prob(self, var_name_log_probs, prediction_target, src_ast_encoding): + def get_target_log_prob( + self, var_name_log_probs, prediction_target, src_ast_encoding + ): # (batch_size, max_prediction_node_num) - variable_tgt_name_id = prediction_target['variable_tgt_name_id'] - tgt_var_name_log_prob = torch.gather(var_name_log_probs, - dim=-1, index=variable_tgt_name_id.unsqueeze(-1)).squeeze(-1) + variable_tgt_name_id = prediction_target["variable_tgt_name_id"] + tgt_var_name_log_prob = torch.gather( + var_name_log_probs, dim=-1, index=variable_tgt_name_id.unsqueeze(-1) + ).squeeze(-1) - tgt_var_name_log_prob = tgt_var_name_log_prob * prediction_target['target_variable_encoding_indices_mask'] + tgt_var_name_log_prob = ( + tgt_var_name_log_prob + * prediction_target["target_variable_encoding_indices_mask"] + ) result = dict(tgt_var_name_log_prob=tgt_var_name_log_prob) @@ -153,7 +196,7 @@ def get_target_log_prob(self, var_name_log_probs, prediction_target, src_ast_enc def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: batch_size = len(examples) - beam_size = self.config['beam_size'] + beam_size = self.config["beam_size"] same_variable_id = self.vocab.target[SAME_VARIABLE_TOKEN] end_of_variable_id = self.vocab.target[END_OF_VARIABLE_TOKEN] @@ -161,7 +204,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: for ast_id, example in enumerate(examples): variable_nums.append(len(example.ast.variables)) - beams = OrderedDict((ast_id, [self.Hypothesis([], 0, 0.)]) for ast_id in range(batch_size)) + beams = OrderedDict( + (ast_id, [self.Hypothesis([], 0, 0.0)]) for ast_id in range(batch_size) + ) hyp_scores_tm1 = torch.zeros(len(beams), device=self.device) completed_hyps = [[] for _ in range(batch_size)] tgt_vocab_size = len(self.vocab.target) @@ -174,20 +219,26 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: # Note that we are using the `restoration_indices` from `context_encoding`, which is the word-level restoration index # (batch_size, variable_master_node_num, encoding_size) - variable_encoding = context_encoding['variable_encoding'] + variable_encoding = context_encoding["variable_encoding"] # (batch_size, encoding_size) - variable_name_embed_tm1 = att_tm1 = torch.zeros(batch_size, self.lstm_cell.hidden_size, device=self.device) + variable_name_embed_tm1 = att_tm1 = torch.zeros( + batch_size, self.lstm_cell.hidden_size, device=self.device + ) - max_prediction_time_step = self.config['max_prediction_time_step'] + max_prediction_time_step = self.config["max_prediction_time_step"] for t in range(0, max_prediction_time_step): # (total_live_hyp_num, encoding_size) if t > 0: - variable_encoding_t = variable_encoding[hyp_ast_ids_t, hyp_variable_ptrs_t] + variable_encoding_t = variable_encoding[ + hyp_ast_ids_t, hyp_variable_ptrs_t + ] else: variable_encoding_t = variable_encoding[:, 0] - if self.config['input_feed']: - x = torch.cat([variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1) + if self.config["input_feed"]: + x = torch.cat( + [variable_encoding_t, variable_name_embed_tm1, att_tm1], dim=-1 + ) else: x = torch.cat([variable_encoding_t, variable_name_embed_tm1], dim=-1) @@ -210,11 +261,13 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: for beam_id, (ast_id, beam) in enumerate(beams.items()): beam_end_hyp_pos = beam_start_hyp_pos + len(beam) # (live_beam_size, vocab_size) - beam_cont_cand_hyp_scores = cont_cand_hyp_scores[beam_start_hyp_pos: beam_end_hyp_pos] + beam_cont_cand_hyp_scores = cont_cand_hyp_scores[ + beam_start_hyp_pos:beam_end_hyp_pos + ] cont_beam_size = beam_size - len(completed_hyps[ast_id]) - beam_new_hyp_scores, beam_new_hyp_positions = torch.topk(beam_cont_cand_hyp_scores.view(-1), - k=cont_beam_size, - dim=-1) + beam_new_hyp_scores, beam_new_hyp_positions = torch.topk( + beam_cont_cand_hyp_scores.view(-1), k=cont_beam_size, dim=-1 + ) # (cont_beam_size) beam_prev_hyp_ids = beam_new_hyp_positions / tgt_vocab_size @@ -235,12 +288,17 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: variable_ptr += 1 # remove empty cases - if len(prev_hyp.variable_list) == 0 or prev_hyp.variable_list[-1] == end_of_variable_id: + if ( + len(prev_hyp.variable_list) == 0 + or prev_hyp.variable_list[-1] == end_of_variable_id + ): continue - new_hyp = self.Hypothesis(variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], - variable_ptr=variable_ptr, - score=new_hyp_score) + new_hyp = self.Hypothesis( + variable_list=list(prev_hyp.variable_list) + [hyp_var_name_id], + variable_ptr=variable_ptr, + score=new_hyp_score, + ) if variable_ptr == variable_nums[ast_id]: completed_hyps[ast_id].append(new_hyp) @@ -252,7 +310,9 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: new_hyp_var_name_ids.append(hyp_var_name_id) new_hyp_ast_ids.append(ast_id) new_hyp_variable_ptrs.append(variable_ptr) - is_same_variable_mask.append(1. if prev_hyp.variable_ptr == variable_ptr else 0.) + is_same_variable_mask.append( + 1.0 if prev_hyp.variable_ptr == variable_ptr else 0.0 + ) beam_start_hyp_pos = beam_end_hyp_pos @@ -268,10 +328,17 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: beams = new_beams if self.independent_prediction_for_each_variable: - is_same_variable_mask = torch.tensor(is_same_variable_mask, device=self.device, dtype=torch.float).unsqueeze(-1) - h_tm1 = (h_tm1[0] * is_same_variable_mask, h_tm1[1] * is_same_variable_mask) + is_same_variable_mask = torch.tensor( + is_same_variable_mask, device=self.device, dtype=torch.float + ).unsqueeze(-1) + h_tm1 = ( + h_tm1[0] * is_same_variable_mask, + h_tm1[1] * is_same_variable_mask, + ) att_tm1 = att_tm1 * is_same_variable_mask - variable_name_embed_tm1 = variable_name_embed_tm1 * is_same_variable_mask + variable_name_embed_tm1 = ( + variable_name_embed_tm1 * is_same_variable_mask + ) else: break @@ -283,10 +350,15 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: if not hyps: # return identity renamings - print(f'Failed to found a hypothesis for function {ast.compilation_unit}', file=sys.stderr) + print( + f"Failed to found a hypothesis for function {ast.compilation_unit}", + file=sys.stderr, + ) for old_name in ast.variables: - variable_rename_result[old_name] = {'new_name': old_name, - 'prob': 0.} + variable_rename_result[old_name] = { + "new_name": old_name, + "prob": 0.0, + } else: top_hyp = hyps[0] sub_token_ptr = 0 @@ -297,14 +369,20 @@ def predict(self, examples: List[Example], encoder: Encoder) -> List[Dict]: sub_token_ptr += 1 # point to first sub-token of next variable sub_token_end = sub_token_ptr - var_name_token_ids = top_hyp.variable_list[sub_token_begin: sub_token_end] # include ending + var_name_token_ids = top_hyp.variable_list[ + sub_token_begin:sub_token_end + ] # include ending if var_name_token_ids == [same_variable_id, end_of_variable_id]: new_var_name = old_name else: - new_var_name = self.vocab.target.subtoken_model.decode_ids(var_name_token_ids) - - variable_rename_result[old_name] = {'new_name': new_var_name, - 'prob': top_hyp.score} + new_var_name = self.vocab.target.subtoken_model.decode_ids( + var_name_token_ids + ) + + variable_rename_result[old_name] = { + "new_name": new_var_name, + "prob": top_hyp.score, + } variable_rename_results.append(variable_rename_result) diff --git a/prediction-plugin/model/sequential_encoder.py b/prediction-plugin/model/sequential_encoder.py index db00f85..a8f0db6 100644 --- a/prediction-plugin/model/sequential_encoder.py +++ b/prediction-plugin/model/sequential_encoder.py @@ -1,31 +1,39 @@ import sys from typing import Dict, List -import numpy as np - -from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence +import numpy as np +import torch +import torch.nn as nn from model.encoder import * +from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence from utils import nn_util, util from utils.dataset import Example from utils.vocab import PAD_ID, Vocab -import torch -import torch.nn as nn class SequentialEncoder(Encoder): def __init__(self, config): super().__init__() - self.vocab = vocab = Vocab.load(config['vocab_file']) + self.vocab = vocab = Vocab.load(config["vocab_file"]) - self.src_word_embed = nn.Embedding(len(vocab.source_tokens), config['source_embedding_size']) + self.src_word_embed = nn.Embedding( + len(vocab.source_tokens), config["source_embedding_size"] + ) - dropout = config['dropout'] - self.lstm_encoder = nn.LSTM(input_size=self.src_word_embed.embedding_dim, - hidden_size=config['source_encoding_size'] // 2, num_layers=config['num_layers'], - batch_first=True, bidirectional=True, dropout=dropout) + dropout = config["dropout"] + self.lstm_encoder = nn.LSTM( + input_size=self.src_word_embed.embedding_dim, + hidden_size=config["source_encoding_size"] // 2, + num_layers=config["num_layers"], + batch_first=True, + bidirectional=True, + dropout=dropout, + ) - self.decoder_cell_init = nn.Linear(config['source_encoding_size'], config['decoder_hidden_size']) + self.decoder_cell_init = nn.Linear( + config["source_encoding_size"], config["decoder_hidden_size"] + ) self.dropout = nn.Dropout(dropout) self.config = config @@ -37,11 +45,11 @@ def device(self): @classmethod def default_params(cls): return { - 'source_encoding_size': 256, - 'decoder_hidden_size': 128, - 'source_embedding_size': 128, - 'vocab_file': None, - 'num_layers': 1 + "source_encoding_size": 256, + "decoder_hidden_size": 128, + "source_embedding_size": 128, + "vocab_file": None, + "num_layers": 1, } @classmethod @@ -51,16 +59,20 @@ def build(cls, config): return cls(params) def forward(self, tensor_dict: Dict[str, torch.Tensor]): - code_token_encoding, code_token_mask, (last_states, last_cells) = self.encode_sequence(tensor_dict['src_code_tokens']) + ( + code_token_encoding, + code_token_mask, + (last_states, last_cells), + ) = self.encode_sequence(tensor_dict["src_code_tokens"]) # (batch_size, max_variable_mention_num) # variable_mention_positions = tensor_dict['variable_position'] - variable_mention_mask = tensor_dict['variable_mention_mask'] - variable_mention_to_variable_id = tensor_dict['variable_mention_to_variable_id'] + variable_mention_mask = tensor_dict["variable_mention_mask"] + variable_mention_to_variable_id = tensor_dict["variable_mention_to_variable_id"] # (batch_size, max_variable_num) - variable_encoding_mask = tensor_dict['variable_encoding_mask'] - variable_mention_num = tensor_dict['variable_mention_num'] + variable_encoding_mask = tensor_dict["variable_encoding_mask"] + variable_mention_num = tensor_dict["variable_mention_num"] # # (batch_size, max_variable_mention_num, encoding_size) # variable_mention_encoding = torch.gather(code_token_encoding, 1, variable_mention_positions.unsqueeze(-1).expand(-1, -1, code_token_encoding.size(-1))) * variable_mention_positions_mask @@ -68,19 +80,27 @@ def forward(self, tensor_dict: Dict[str, torch.Tensor]): variable_num = variable_mention_num.size(1) encoding_size = code_token_encoding.size(-1) - variable_mention_encoding = code_token_encoding * variable_mention_mask.unsqueeze(-1) - variable_encoding = torch.zeros(tensor_dict['batch_size'], variable_num, encoding_size, device=self.device) - variable_encoding.scatter_add_(1, - variable_mention_to_variable_id.unsqueeze(-1).expand(-1, -1, encoding_size), - variable_mention_encoding) * variable_encoding_mask.unsqueeze(-1) - variable_encoding = variable_encoding / (variable_mention_num + (1. - variable_encoding_mask) * nn_util.SMALL_NUMBER).unsqueeze(-1) + variable_mention_encoding = ( + code_token_encoding * variable_mention_mask.unsqueeze(-1) + ) + variable_encoding = torch.zeros( + tensor_dict["batch_size"], variable_num, encoding_size, device=self.device + ) + variable_encoding.scatter_add_( + 1, + variable_mention_to_variable_id.unsqueeze(-1).expand(-1, -1, encoding_size), + variable_mention_encoding, + ) * variable_encoding_mask.unsqueeze(-1) + variable_encoding = variable_encoding / ( + variable_mention_num + (1.0 - variable_encoding_mask) * nn_util.SMALL_NUMBER + ).unsqueeze(-1) context_encoding = dict( variable_encoding=variable_encoding, code_token_encoding=code_token_encoding, code_token_mask=code_token_mask, last_states=last_states, - last_cells=last_cells + last_cells=last_cells, ) context_encoding.update(tensor_dict) @@ -99,12 +119,20 @@ def encode_sequence(self, code_sequence): # (batch_size) code_sequence_length = code_token_mask.sum(dim=-1).long() - sorted_seqs, sorted_seq_lens, restoration_indices, sorting_indices = nn_util.sort_batch_by_length(code_token_embedding, - code_sequence_length) + ( + sorted_seqs, + sorted_seq_lens, + restoration_indices, + sorting_indices, + ) = nn_util.sort_batch_by_length(code_token_embedding, code_sequence_length) - packed_question_embedding = pack_padded_sequence(sorted_seqs, sorted_seq_lens.data.tolist(), batch_first=True) + packed_question_embedding = pack_padded_sequence( + sorted_seqs, sorted_seq_lens.data.tolist(), batch_first=True + ) - sorted_encodings, (last_states, last_cells) = self.lstm_encoder(packed_question_embedding) + sorted_encodings, (last_states, last_cells) = self.lstm_encoder( + packed_question_embedding + ) sorted_encodings, _ = pad_packed_sequence(sorted_encodings, batch_first=True) # apply dropout to the last layer @@ -112,12 +140,18 @@ def encode_sequence(self, code_sequence): sorted_encodings = self.dropout(sorted_encodings) # (batch_size, question_len, hidden_size * 2) - restored_encodings = sorted_encodings.index_select(dim=0, index=restoration_indices) + restored_encodings = sorted_encodings.index_select( + dim=0, index=restoration_indices + ) # (num_layers, direction_num, batch_size, hidden_size) - last_states = last_states.view(self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size) + last_states = last_states.view( + self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size + ) last_states = last_states.index_select(dim=2, index=restoration_indices) - last_cells = last_cells.view(self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size) + last_cells = last_cells.view( + self.lstm_encoder.num_layers, 2, -1, self.lstm_encoder.hidden_size + ) last_cells = last_cells.index_select(dim=2, index=restoration_indices) return restored_encodings, code_token_mask, (last_states, last_cells) @@ -127,25 +161,31 @@ def to_tensor_dict(cls, examples: List[Example]) -> Dict[str, torch.Tensor]: max_time_step = max(e.source_seq_length for e in examples) input = np.zeros((len(examples), max_time_step), dtype=np.int64) - variable_mention_to_variable_id = torch.zeros(len(examples), max_time_step, dtype=torch.long) + variable_mention_to_variable_id = torch.zeros( + len(examples), max_time_step, dtype=torch.long + ) variable_mention_mask = torch.zeros(len(examples), max_time_step) - variable_mention_num = torch.zeros(len(examples), max(len(e.ast.variables) for e in examples)) + variable_mention_num = torch.zeros( + len(examples), max(len(e.ast.variables) for e in examples) + ) variable_encoding_mask = torch.zeros(variable_mention_num.size()) for e_id, example in enumerate(examples): sub_tokens = example.sub_tokens - input[e_id, :len(sub_tokens)] = example.sub_token_ids + input[e_id, : len(sub_tokens)] = example.sub_token_ids variable_position_map = dict() var_name_to_id = {name: i for i, name in enumerate(example.ast.variables)} for i, sub_token in enumerate(sub_tokens): - if sub_token.startswith('@@') and sub_token.endswith('@@'): - old_var_name = sub_token[2: -2] - if old_var_name in var_name_to_id: # sometimes there are strings like `@@@@` + if sub_token.startswith("@@") and sub_token.endswith("@@"): + old_var_name = sub_token[2:-2] + if ( + old_var_name in var_name_to_id + ): # sometimes there are strings like `@@@@` var_id = var_name_to_id[old_var_name] variable_mention_to_variable_id[e_id, i] = var_id - variable_mention_mask[e_id, i] = 1. + variable_mention_mask[e_id, i] = 1.0 variable_position_map.setdefault(old_var_name, []).append(i) for var_id, var_name in enumerate(example.ast.variables): @@ -154,30 +194,38 @@ def to_tensor_dict(cls, examples: List[Example]) -> Dict[str, torch.Tensor]: variable_mention_num[e_id, var_id] = len(var_pos) except KeyError: variable_mention_num[e_id, var_id] = 1 - print(example.binary_file, f'variable [{var_name}] not found', file=sys.stderr) - - variable_encoding_mask[e_id, :len(example.ast.variables)] = 1. - - return dict(src_code_tokens=torch.from_numpy(input), - variable_mention_to_variable_id=variable_mention_to_variable_id, - variable_mention_mask=variable_mention_mask, - variable_mention_num=variable_mention_num, - variable_encoding_mask=variable_encoding_mask, - batch_size=len(examples)) + print( + example.binary_file, + f"variable [{var_name}] not found", + file=sys.stderr, + ) + + variable_encoding_mask[e_id, : len(example.ast.variables)] = 1.0 + + return dict( + src_code_tokens=torch.from_numpy(input), + variable_mention_to_variable_id=variable_mention_to_variable_id, + variable_mention_mask=variable_mention_mask, + variable_mention_num=variable_mention_num, + variable_encoding_mask=variable_encoding_mask, + batch_size=len(examples), + ) def get_decoder_init_state(self, context_encoder, config=None): - fwd_last_layer_cell = context_encoder['last_cells'][-1, 0] - bak_last_layer_cell = context_encoder['last_cells'][-1, 1] + fwd_last_layer_cell = context_encoder["last_cells"][-1, 0] + bak_last_layer_cell = context_encoder["last_cells"][-1, 1] - dec_init_cell = self.decoder_cell_init(torch.cat([fwd_last_layer_cell, bak_last_layer_cell], dim=-1)) + dec_init_cell = self.decoder_cell_init( + torch.cat([fwd_last_layer_cell, bak_last_layer_cell], dim=-1) + ) dec_init_state = torch.tanh(dec_init_cell) return dec_init_state, dec_init_cell - def get_attention_memory(self, context_encoding, att_target='terminal_nodes'): - assert att_target == 'terminal_nodes' + def get_attention_memory(self, context_encoding, att_target="terminal_nodes"): + assert att_target == "terminal_nodes" - memory = context_encoding['code_token_encoding'] - mask = context_encoding['code_token_mask'] + memory = context_encoding["code_token_encoding"] + mask = context_encoding["code_token_mask"] return memory, mask diff --git a/prediction-plugin/model/simple_decoder.py b/prediction-plugin/model/simple_decoder.py new file mode 100644 index 0000000..127ecaa --- /dev/null +++ b/prediction-plugin/model/simple_decoder.py @@ -0,0 +1,135 @@ +from typing import Dict, List + +import torch +from model.decoder import Decoder +from torch import nn as nn +from utils import nn_util, util +from utils.ast import AbstractSyntaxTree +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab + + +class SimpleDecoder(Decoder): + def __init__(self, ast_node_encoding_size: int, vocab: Vocab): + super(SimpleDecoder, self).__init__() + + self.vocab = vocab + self.state2names = nn.Linear( + ast_node_encoding_size, len(vocab.target), bias=True + ) + self.config: Dict = None + + @classmethod + def default_params(cls): + return {"vocab_file": None, "ast_node_encoding_size": 128} + + @classmethod + def build(cls, config): + params = util.update(cls.default_params(), config) + + vocab = torch.load(params["vocab_file"]) + model = cls(params["ast_node_encoding_size"], vocab) + model.config = params + + return model + + def forward( + self, + src_ast_encoding: Dict[str, torch.Tensor], + prediction_target: Dict[str, torch.Tensor], + ): + """ + Given a batch of encoded ASTs, compute the log-likelihood of generating all possible renamings + """ + # (all_var_node_num, tgt_vocab_size) + logits = self.state2names(src_ast_encoding["variable_master_node_encoding"]) + batched_p_names = torch.log_softmax(logits, dim=-1) + # logits = self.state2names(src_ast_encoding) + # p = torch.log_softmax(logits, dim=-1) + + # idx = src_ast_encoding.unpacked_variable_node_ids + # (batch_size, max_variable_node_num, tgt_vocab_size) + # batched_p_names.unsqueeze(-1).expand_as(src_ast_encoding.batch_size, -1, -1).scatter(idx, dim=1) + + # (batch_var_node_num) + packed_tgt_var_node_name_log_probs = torch.gather( + batched_p_names, + dim=-1, + index=prediction_target["variable_tgt_name_id"].unsqueeze(-1), + ).squeeze(-1) + + # result = dict(context_encoding=context_encoding) + result = dict() + with torch.no_grad(): + # compute ppl over renamed variables + renamed_var_mask = prediction_target["var_with_new_name_mask"] + unchanged_var_mask = prediction_target["auxiliary_var_mask"] + renamed_var_avg_ll = ( + packed_tgt_var_node_name_log_probs * renamed_var_mask + ).sum(-1) / renamed_var_mask.sum() + unchanged_var_avg_ll = ( + packed_tgt_var_node_name_log_probs * unchanged_var_mask + ).sum(-1) / unchanged_var_mask.sum() + result["rename_ppl"] = torch.exp(-renamed_var_avg_ll).item() + result["unchange_ppl"] = torch.exp(-unchanged_var_avg_ll).item() + + packed_tgt_var_node_name_log_probs = ( + packed_tgt_var_node_name_log_probs + * prediction_target["variable_tgt_name_weight"] + ) + + # (batch_size, max_variable_node_num) + tgt_name_log_probs = packed_tgt_var_node_name_log_probs[ + src_ast_encoding["variable_encoding_restoration_indices"] + ] + tgt_name_log_probs = ( + tgt_name_log_probs + * src_ast_encoding["variable_encoding_restoration_indices_mask"] + ) + + # (batch_size) + ast_log_probs = tgt_name_log_probs.sum(dim=-1) / src_ast_encoding[ + "variable_encoding_restoration_indices_mask" + ].sum(-1) + + result["batch_log_prob"] = ast_log_probs + + return result + + def predict(self, source_asts: List[AbstractSyntaxTree]): + """ + Given a batch of ASTs, predict their new variable names + """ + + tensor_dict = self.batcher.to_tensor_dict(source_asts=source_asts) + nn_util.to(tensor_dict, self.device) + context_encoding = self.encoder(tensor_dict) + # (prediction_size, tgt_vocab_size) + packed_var_name_log_probs = self.decoder(context_encoding) + best_var_name_log_probs, best_var_name_ids = torch.max( + packed_var_name_log_probs, dim=-1 + ) + + variable_rename_results = [] + pred_node_ptr = 0 + for ast_id, ast in enumerate(source_asts): + variable_rename_result = dict() + for var_name in ast.variables: + var_name_prob = best_var_name_log_probs[pred_node_ptr].item() + token_id = best_var_name_ids[pred_node_ptr].item() + new_var_name = self.vocab.target.id2word[token_id] + + if new_var_name == SAME_VARIABLE_TOKEN: + new_var_name = var_name + + variable_rename_result[var_name] = { + "new_name": new_var_name, + "prob": var_name_prob, + } + + pred_node_ptr += 1 + + variable_rename_results.append(variable_rename_result) + + assert pred_node_ptr == packed_var_name_log_probs.size(0) + + return variable_rename_results diff --git a/prediction-plugin/utils/ast.py b/prediction-plugin/utils/ast.py index 0492c97..766789f 100644 --- a/prediction-plugin/utils/ast.py +++ b/prediction-plugin/utils/ast.py @@ -1,19 +1,23 @@ from collections import OrderedDict from io import StringIO -import ujson as json from typing import Dict, List +import ujson as json from utils.util import cached_property +from utils.vocab import VocabEntry class SyntaxNode(object): """represent a node on an AST""" - def __init__(self, - node_id, - node_type, - address=None, - children: List = None, - named_fields: Dict = None): + + def __init__( + self, + node_id, + node_type, + address=None, + children: List = None, + named_fields: Dict = None, + ): self.node_id = node_id self.node_type = node_type self.address = address @@ -30,33 +34,35 @@ def __init__(self, for child in children: self.add_child(child) - def add_child(self, child: 'SyntaxNode') -> None: + def add_child(self, child: "SyntaxNode") -> None: self.children.append(child) child.parent = self @classmethod - def from_json_dict(cls, json_dict: Dict) -> 'SyntaxNode': + def from_json_dict(cls, json_dict: Dict) -> "SyntaxNode": named_fields = { - k: v for k, v in json_dict.items() - if k not in {'node_id', 'node_type', 'children', - 'address', 'x', 'y', 'z'} + k: v + for k, v in json_dict.items() + if k not in {"node_id", "node_type", "children", "address", "x", "y", "z"} } - if 'x' in json_dict: - named_fields['x'] = SyntaxNode.from_json_dict(json_dict['x']) - if 'y' in json_dict: - named_fields['y'] = SyntaxNode.from_json_dict(json_dict['y']) - if 'z' in json_dict: - named_fields['z'] = SyntaxNode.from_json_dict(json_dict['z']) - - node = cls(json_dict['node_id'], - json_dict['node_type'], - json_dict['address'], - named_fields=named_fields) + if "x" in json_dict: + named_fields["x"] = SyntaxNode.from_json_dict(json_dict["x"]) + if "y" in json_dict: + named_fields["y"] = SyntaxNode.from_json_dict(json_dict["y"]) + if "z" in json_dict: + named_fields["z"] = SyntaxNode.from_json_dict(json_dict["z"]) + + node = cls( + json_dict["node_id"], + json_dict["node_type"], + json_dict["address"], + named_fields=named_fields, + ) children_list = [] - if 'children' in json_dict: - children_list.extend(json_dict['children']) + if "children" in json_dict: + children_list.extend(json_dict["children"]) for child_dict in children_list: child_node = SyntaxNode.from_json_dict(child_dict) @@ -65,13 +71,13 @@ def from_json_dict(cls, json_dict: Dict) -> 'SyntaxNode': return node def to_json_dict(self): - json_dict = dict(node_id=self.node_id, - node_type=self.node_type, - address=self.address) + json_dict = dict( + node_id=self.node_id, node_type=self.node_type, address=self.address + ) for named_filed in self.named_fields: val = getattr(self, named_filed) - if named_filed in ('x', 'y', 'z'): + if named_filed in ("x", "y", "z"): json_dict[named_filed] = val.to_json_dict() else: json_dict[named_filed] = val @@ -80,20 +86,22 @@ def to_json_dict(self): children = [] for child in self.children: children.append(child.to_json_dict()) - json_dict['children'] = children + json_dict["children"] = children return json_dict @property def is_variable_node(self): - return self.node_type == 'var' + return self.node_type == "var" @property def is_terminal_node(self): - return not hasattr(self, 'x') \ - and not hasattr(self, 'y') \ - and not hasattr(self, 'z') \ + return ( + not hasattr(self, "x") + and not hasattr(self, "y") + and not hasattr(self, "z") and not self.children + ) @cached_property def size(self): @@ -105,11 +113,11 @@ def size(self): @property def member_nodes(self): - if hasattr(self, 'x'): + if hasattr(self, "x"): yield self.x - if hasattr(self, 'y'): + if hasattr(self, "y"): yield self.y - if hasattr(self, 'z'): + if hasattr(self, "z"): yield self.z for child in self.children: @@ -117,14 +125,14 @@ def member_nodes(self): @property def named_succeeding_fields(self): - if hasattr(self, 'x'): - yield 'x', self.x - if hasattr(self, 'y'): - yield 'y', self.y - if hasattr(self, 'z'): - yield 'z', self.z + if hasattr(self, "x"): + yield "x", self.x + if hasattr(self, "y"): + yield "y", self.y + if hasattr(self, "z"): + yield "z", self.z - yield 'children', self.children + yield "children", self.children @property def descendant_nodes(self): @@ -133,6 +141,7 @@ def _visit(node): for member_node in node.member_nodes: yield from _visit(member_node) + yield from _visit(self) def __iter__(self): @@ -146,17 +155,20 @@ def __hash__(self): return code def __eq__(self, other): - if not isinstance(other, self.__class__) \ - or self.node_type != other.node_type \ - or self.address != other.address \ - or self.node_id != other.node_id \ - or self.named_fields != other.named_fields \ - or len(self.children) != len(other.children): + if ( + not isinstance(other, self.__class__) + or self.node_type != other.node_type + or self.address != other.address + or self.node_id != other.node_id + or self.named_fields != other.named_fields + or len(self.children) != len(other.children) + ): return False for i in range(len(self.children)): if self.children[i] != other.children[i]: return False + return True def __ne__(self, other): @@ -168,38 +180,35 @@ def to_string(self, sb=None): is_root = True sb = StringIO() - sb.write(f'(Node{self.node_id}-{self.node_type}') + sb.write(f"(Node{self.node_id}-{self.node_type}") for key in self.named_fields: val = getattr(self, key) - if key not in ('x', 'y', 'z'): - sb.write('-') + if key not in ("x", "y", "z"): + sb.write("-") sb.write( - f'{key}:{val}' - .replace(' ', '_') - .replace('(', '_') - .replace(')', '_') + f"{key}:{val}".replace(" ", "_").replace("(", "_").replace(")", "_") ) # sb.write(f'Node-{self.node_id}-{self.node_type}') for field_name, node in self.named_succeeding_fields: - if field_name in ('x', 'y'): - sb.write(f' ({field_name} ') + if field_name in ("x", "y"): + sb.write(f" ({field_name} ") node.to_string(sb) - sb.write(')') # of x + sb.write(")") # of x elif self.children: - sb.write(' (children ') + sb.write(" (children ") for child in self.children: - sb.write(' ') + sb.write(" ") child.to_string(sb) - sb.write(')') # of children field + sb.write(")") # of children field - sb.write(')') # of node + sb.write(")") # of node if is_root: return sb.getvalue() def __str__(self): - return f'Node {self.node_id} {self.node_type}@{self.address}' + return f"Node {self.node_id} {self.node_type}@{self.address}" __repr__ = __str__ @@ -208,14 +217,14 @@ class TerminalNode(SyntaxNode): """a terminal AST node representing variables or other terminal syntax tokens """ + pass class AbstractSyntaxTree(object): - def __init__(self, - root: SyntaxNode, - compilation_unit: str = None, - code: str = None): + def __init__( + self, root: SyntaxNode, compilation_unit: str = None, code: str = None + ): self.root = root self.compilation_unit = compilation_unit self.code = code @@ -229,13 +238,13 @@ def __init__(self, self._init_index() @classmethod - def from_json_dict(cls, json_dict: Dict) -> 'AbstractSyntaxTree': - root = SyntaxNode.from_json_dict(json_dict['ast']) - root.name = json_dict['function'] + def from_json_dict(cls, json_dict: Dict) -> "AbstractSyntaxTree": + root = SyntaxNode.from_json_dict(json_dict["ast"]) + root.name = json_dict["function"] tree = cls( root, - compilation_unit=json_dict['function'], - code=json_dict.get('raw_code', None) + compilation_unit=json_dict["function"], + code=json_dict.get("raw_code", None), ) return tree @@ -268,21 +277,21 @@ def _index_sub_tree(node: SyntaxNode, parent_node: SyntaxNode): _index_sub_tree(self.root, None) - setattr(self, 'adjacency_list', adj_list) - setattr(self, 'id_to_node', id2node) + setattr(self, "adjacency_list", adj_list) + setattr(self, "id_to_node", id2node) # TODO: implement this! - setattr(self, 'adjacent_terminal_nodes', []) - setattr(self, 'variable_nodes', variable_nodes) - setattr(self, 'variables', variables) + setattr(self, "adjacent_terminal_nodes", []) + setattr(self, "variable_nodes", variable_nodes) + setattr(self, "variables", variables) # TODO: change to address based! terminal_nodes.sort(key=lambda n: n.node_id) - setattr(self, 'terminal_nodes', terminal_nodes) + setattr(self, "terminal_nodes", terminal_nodes) def __iter__(self): return iter((node for node in self.id_to_node.values())) -if __name__ == '__main__': +if __name__ == "__main__": json_str = """ { "function": "ft_strncat", @@ -715,7 +724,7 @@ def __iter__(self): } """ json_dict = json.loads(json_str) - tree = SyntaxNode.from_json_dict(json_dict['root']) + tree = SyntaxNode.from_json_dict(json_dict["root"]) tree_reconstr = SyntaxNode.from_json_dict(tree.to_json_dict()) assert tree_reconstr == tree diff --git a/prediction-plugin/utils/code_processing.py b/prediction-plugin/utils/code_processing.py index bc6ca79..0cf2775 100644 --- a/prediction-plugin/utils/code_processing.py +++ b/prediction-plugin/utils/code_processing.py @@ -4,27 +4,26 @@ from utils.ast import SyntaxNode from utils.lexer import Lexer, Token - -VARIABLE_ANNOTATION = re.compile(r'@@\w+@@(\w+)@@\w+') +VARIABLE_ANNOTATION = re.compile(r"@@\w+@@(\w+)@@\w+") def canonicalize_code(code): - code = re.sub('//.*?\\n|/\\*.*?\\*/', '\\n', code, flags=re.S) - lines = [l.rstrip() for l in code.split('\\n')] - code = '\\n'.join(lines) - code = re.sub('@@\\w+@@(\\w+)@@\\w+', '\\g<1>', code) + code = re.sub("//.*?\\n|/\\*.*?\\*/", "\\n", code, flags=re.S) + lines = [l.rstrip() for l in code.split("\\n")] + code = "\\n".join(lines) + code = re.sub("@@\\w+@@(\\w+)@@\\w+", "\\g<1>", code) return code def canonicalize_constants(root: SyntaxNode) -> None: def _visit(node): - if node.node_type == 'obj' and node.type == 'char *': - node.name = 'STRING' - elif node.node_type == 'num': - node.name = 'NUMBER' - elif node.node_type == 'fnum': - node.name = 'FLOAT' + if node.node_type == "obj" and node.type == "char *": + node.name = "STRING" + elif node.node_type == "num": + node.name = "NUMBER" + elif node.node_type == "fnum": + node.name = "FLOAT" for child in node.member_nodes: _visit(child) @@ -34,12 +33,11 @@ def _visit(node): def annotate_type(root: SyntaxNode) -> None: def _visit(node): - if hasattr(node, 'type'): - type_tokens = [t[1].lstrip('_') - for t in Lexer(node.type).get_tokens()] - type_tokens = [t for t in type_tokens if t not in ('(', ')')] - node.named_fields.add('type_tokens') - setattr(node, 'type_tokens', type_tokens) + if hasattr(node, "type"): + type_tokens = [t[1].lstrip("_") for t in Lexer(node.type).get_tokens()] + type_tokens = [t for t in type_tokens if t not in ("(", ")")] + node.named_fields.add("type_tokens") + setattr(node, "type_tokens", type_tokens) for child in node.member_nodes: _visit(child) @@ -50,42 +48,37 @@ def _visit(node): VAR_ID_REGEX = re.compile(r"@@(VAR_\d+)@@") -def preprocess_ast(root: SyntaxNode, - preprocessors: Set[str] = None, - code: str = None) -> None: +def preprocess_ast( + root: SyntaxNode, preprocessors: Set[str] = None, code: str = None +) -> None: if preprocessors is None: - preprocessors = { - 'annotate_type', - 'canonicalize_constant', - 'annotate_arg' - } + preprocessors = {"annotate_type", "canonicalize_constant", "annotate_arg"} arg_var_ids = None - if 'annotate_arg' in preprocessors: - first_line = code[:code.index('\n')] + if "annotate_arg" in preprocessors: + first_line = code[: code.index("\n")] arg_var_ids = set(VAR_ID_REGEX.findall(first_line)) def _visit(node): - if 'annotate_type' in preprocessors: - if node.node_type == 'obj' and node.type == 'char *': - node.name = 'STRING' - elif node.node_type == 'num': - node.name = 'NUMBER' - elif node.node_type == 'fnum': - node.name = 'FLOAT' - - if 'canonicalize_constant' in preprocessors: - if hasattr(node, 'type'): - type_tokens = [t[1].lstrip('_') - for t in Lexer(node.type).get_tokens()] - type_tokens = [t for t in type_tokens if t not in ('(', ')')] - node.named_fields.add('type_tokens') - setattr(node, 'type_tokens', type_tokens) - - if 'annotate_arg' in preprocessors: - if node.node_type == 'var': - node.named_fields.add('is_arg') - setattr(node, 'is_arg', node.var_id in arg_var_ids) + if "annotate_type" in preprocessors: + if node.node_type == "obj" and node.type == "char *": + node.name = "STRING" + elif node.node_type == "num": + node.name = "NUMBER" + elif node.node_type == "fnum": + node.name = "FLOAT" + + if "canonicalize_constant" in preprocessors: + if hasattr(node, "type"): + type_tokens = [t[1].lstrip("_") for t in Lexer(node.type).get_tokens()] + type_tokens = [t for t in type_tokens if t not in ("(", ")")] + node.named_fields.add("type_tokens") + setattr(node, "type_tokens", type_tokens) + + if "annotate_arg" in preprocessors: + if node.node_type == "var": + node.named_fields.add("is_arg") + setattr(node, "is_arg", node.var_id in arg_var_ids) for child in node.member_nodes: _visit(child) @@ -98,12 +91,12 @@ def tokenize_raw_code(raw_code): tokens = [] for token_type, token in lexer.get_tokens(): if token_type in Token.Literal: - token = str(token_type).split('.')[2] + token = str(token_type).split(".")[2] if token_type == Token.Placeholder.Var: m = VARIABLE_ANNOTATION.match(token) old_name = m.group(1) - token = '@@' + old_name + '@@' + token = "@@" + old_name + "@@" tokens.append(token) diff --git a/prediction-plugin/utils/dataset.py b/prediction-plugin/utils/dataset.py index ece988b..ad535e7 100644 --- a/prediction-plugin/utils/dataset.py +++ b/prediction-plugin/utils/dataset.py @@ -2,37 +2,37 @@ import glob import multiprocessing import os +import pickle import queue +import random import resource import sys import tarfile import threading import time -import torch +from typing import Dict, Iterable, List, Tuple, Union import numpy as np +import sentencepiece as spm +import torch import torch.multiprocessing as torch_mp import ujson as json - -import random from tqdm import tqdm -from typing import Iterable, List, Dict, Union from utils import nn_util -from utils.ast import AbstractSyntaxTree -from utils.vocab import SAME_VARIABLE_TOKEN, Vocab - +from utils.ast import AbstractSyntaxTree, SyntaxNode +from utils.code_processing import annotate_type +from utils.graph import PackedGraph +from utils.vocab import SAME_VARIABLE_TOKEN, Vocab, VocabEntry batcher_sync_msg = None -torch.multiprocessing.set_sharing_strategy('file_system') + +torch.multiprocessing.set_sharing_strategy("file_system") rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1])) class Example(object): - def __init__(self, - ast: AbstractSyntaxTree, - variable_name_map: dict, - **kwargs): + def __init__(self, ast: AbstractSyntaxTree, variable_name_map: dict, **kwargs): self.ast = ast self.variable_name_map = variable_name_map @@ -47,14 +47,14 @@ def from_json_dict(cls, json_dict, **kwargs): for var_name, var_nodes in tree.variables.items(): variable_name_map[var_name] = var_nodes[0].new_name - if 'test_meta' in json_dict: - kwargs['test_meta'] = json_dict['test_meta'] + if "test_meta" in json_dict: + kwargs["test_meta"] = json_dict["test_meta"] return cls(tree, variable_name_map, **kwargs) class Batch(object): - __slots__ = ('examples', 'tensor_dict') + __slots__ = ("examples", "tensor_dict") def __init__(self, examples, tensor_dict): self.examples = examples @@ -62,7 +62,7 @@ def __init__(self, examples, tensor_dict): @property def size(self): - return self.tensor_dict['batch_size'] + return self.tensor_dict["batch_size"] class Batcher(object): @@ -71,22 +71,25 @@ def __init__(self, config, train=True): self.train = train # model specific config - self.is_ensemble = config['encoder']['type'] == 'EnsembleModel' + self.is_ensemble = config["encoder"]["type"] == "EnsembleModel" if not self.is_ensemble: - self.vocab = Vocab.load(config['data']['vocab_file']) + self.vocab = Vocab.load(config["data"]["vocab_file"]) self.grammar = self.vocab.grammar - self.use_seq_encoder = config['encoder']['type'] == 'SequentialEncoder' - self.use_hybrid_encoder = config['encoder']['type'] == 'HybridEncoder' - self.init_gnn_with_seq_encoding = \ - config['encoder']['type'] == 'GraphASTEncoder' \ - and config['encoder']['init_with_seq_encoding'] + self.use_seq_encoder = config["encoder"]["type"] == "SequentialEncoder" + self.use_hybrid_encoder = config["encoder"]["type"] == "HybridEncoder" + self.init_gnn_with_seq_encoding = ( + config["encoder"]["type"] == "GraphASTEncoder" + and config["encoder"]["init_with_seq_encoding"] + ) @property def annotate_sequential_input(self): - return self.use_seq_encoder \ - or self.use_hybrid_encoder \ + return ( + self.use_seq_encoder + or self.use_hybrid_encoder or self.init_gnn_with_seq_encoding + ) def annotate_example(self, example) -> Example: """Annotate examples by populating specific fields, useful for sorting @@ -100,31 +103,31 @@ def annotate_example(self, example) -> Example: if self.annotate_sequential_input: src_bpe_model = self.vocab.source_tokens.subtoken_model snippet = example.code_tokens - snippet = ' '.join(snippet) - sub_tokens = \ - [''] + src_bpe_model.encode_as_pieces(snippet) + [''] - sub_token_ids = \ - [src_bpe_model.bos_id()] \ - + src_bpe_model.encode_as_ids(snippet) \ + snippet = " ".join(snippet) + sub_tokens = [""] + src_bpe_model.encode_as_pieces(snippet) + [""] + sub_token_ids = ( + [src_bpe_model.bos_id()] + + src_bpe_model.encode_as_ids(snippet) + [src_bpe_model.eos_id()] - setattr(example, 'sub_tokens', sub_tokens) - setattr(example, 'sub_token_ids', sub_token_ids) - setattr(example, 'source_seq_length', len(sub_tokens)) + ) + setattr(example, "sub_tokens", sub_tokens) + setattr(example, "sub_token_ids", sub_token_ids) + setattr(example, "source_seq_length", len(sub_tokens)) tgt_bpe_model = self.vocab.target.subtoken_model eov_id = tgt_bpe_model.eos_id() - var_name_subtoken_map = dict() + variable_name_subtoken_map = dict() tgt_pred_seq_len = 0 for old_name, new_name in example.variable_name_map.items(): if old_name == new_name: subtoken_ids = [self.vocab.target[SAME_VARIABLE_TOKEN], eov_id] else: subtoken_ids = tgt_bpe_model.encode_as_ids(new_name) + [eov_id] - var_name_subtoken_map[old_name] = subtoken_ids + variable_name_subtoken_map[old_name] = subtoken_ids tgt_pred_seq_len += len(subtoken_ids) - setattr(example, 'variable_name_subtoken_map', var_name_subtoken_map) - setattr(example, 'target_prediction_seq_length', tgt_pred_seq_len) + setattr(example, "variable_name_subtoken_map", variable_name_subtoken_map) + setattr(example, "target_prediction_seq_length", tgt_pred_seq_len) return example @@ -148,83 +151,81 @@ def get_batch_size(self, examples: List[Example]): if self.annotate_sequential_input: return len(examples) * max(e.source_seq_length for e in examples) else: - return len(examples) * \ - max(e.target_prediction_seq_length for e in examples) + return len(examples) * max(e.target_prediction_seq_length for e in examples) - def to_tensor_dict(self, - examples: List[Example], - return_prediction_target=True): - # type: (...) -> Dict[str, torch.Tensor] - from model.sequential_encoder import SequentialEncoder + def to_tensor_dict( + self, examples: List[Example], return_prediction_target=True + ) -> Dict[str, torch.Tensor]: from model.graph_encoder import GraphASTEncoder + from model.sequential_encoder import SequentialEncoder - if not hasattr(examples[0], 'target_prediction_seq_length'): + if not hasattr(examples[0], "target_prediction_seq_length"): for example in examples: self.annotate_example(example) - if self.config['encoder']['type'] == 'GraphASTEncoder': - init_with_seq_encoding = \ - self.config['encoder']['init_with_seq_encoding'] - connections = self.config['encoder']['connections'] - packed_graph, tensor_dict = \ - GraphASTEncoder.to_packed_graph( - [e.ast for e in examples], - connections=connections, - init_with_seq_encoding=init_with_seq_encoding - ) + if self.config["encoder"]["type"] == "GraphASTEncoder": + init_with_seq_encoding = self.config["encoder"]["init_with_seq_encoding"] + packed_graph, tensor_dict = GraphASTEncoder.to_packed_graph( + [e.ast for e in examples], + connections=self.config["encoder"]["connections"], + init_with_seq_encoding=init_with_seq_encoding, + ) if init_with_seq_encoding: seq_tensor_dict = SequentialEncoder.to_tensor_dict(examples) - tensor_dict['seq_encoder_input'] = seq_tensor_dict + tensor_dict["seq_encoder_input"] = seq_tensor_dict - _tensors = GraphASTEncoder.to_tensor_dict(packed_graph, - self.grammar, self.vocab) + _tensors = GraphASTEncoder.to_tensor_dict( + packed_graph, self.grammar, self.vocab + ) tensor_dict.update(_tensors) - elif self.config['encoder']['type'] == 'SequentialEncoder': + elif self.config["encoder"]["type"] == "SequentialEncoder": tensor_dict = SequentialEncoder.to_tensor_dict(examples) - elif self.config['encoder']['type'] == 'HybridEncoder': - connections = \ - self.config['encoder']['graph_encoder']['connections'] - packed_graph, gnn_tensor_dict = \ - GraphASTEncoder.to_packed_graph([e.ast for e in examples], - connections=connections) - gnn_tensors = GraphASTEncoder.to_tensor_dict(packed_graph, - self.grammar, - self.vocab) + elif self.config["encoder"]["type"] == "HybridEncoder": + packed_graph, gnn_tensor_dict = GraphASTEncoder.to_packed_graph( + [e.ast for e in examples], + connections=self.config["encoder"]["graph_encoder"]["connections"], + ) + gnn_tensors = GraphASTEncoder.to_tensor_dict( + packed_graph, self.grammar, self.vocab + ) gnn_tensor_dict.update(gnn_tensors) seq_tensor_dict = SequentialEncoder.to_tensor_dict(examples) - tensor_dict = {'graph_encoder_input': gnn_tensor_dict, - 'seq_encoder_input': seq_tensor_dict} + tensor_dict = { + "graph_encoder_input": gnn_tensor_dict, + "seq_encoder_input": seq_tensor_dict, + } else: - raise ValueError('UnknownEncoderType') + raise ValueError("UnknownEncoderType") if self.train or return_prediction_target: prediction_target = self.to_batched_prediction_target(examples) - tensor_dict['prediction_target'] = prediction_target + tensor_dict["prediction_target"] = prediction_target if not self.train: - if hasattr(examples[0], 'test_meta'): - tensor_dict['test_meta'] = [e.test_meta for e in examples] + if hasattr(examples[0], "test_meta"): + tensor_dict["test_meta"] = [e.test_meta for e in examples] - tensor_dict['batch_size'] = len(examples) + tensor_dict["batch_size"] = len(examples) num_elements = nn_util.get_tensor_dict_size(tensor_dict) - tensor_dict['num_elements'] = num_elements + tensor_dict["num_elements"] = num_elements return tensor_dict - def to_batch(self, - examples: List[Example], - return_examples=False, - return_prediction_target=True) -> Batch: + def to_batch( + self, + examples: List[Example], + return_examples=False, + return_prediction_target=True, + ) -> Batch: if self.is_ensemble: # do not perform tensorization for the parent ensemble model tensor_dict = None else: with torch.no_grad(): - tensor_dict = \ - self.to_tensor_dict(examples, return_prediction_target) + tensor_dict = self.to_tensor_dict(examples, return_prediction_target) if not return_examples: batch = Batch(None, tensor_dict) @@ -236,14 +237,13 @@ def to_batch(self, def to_batched_prediction_target(self, examples: List[Example]): batch_size = len(examples) - unchanged_var_weight = \ - self.config['train']['unchanged_variable_weight'] + unchanged_var_weight = self.config["train"]["unchanged_variable_weight"] + use_bpe_for_var_name = self.vocab.target.subtoken_model is not None - var_name_subtoken_maps = [] + variable_name_subtoken_maps = [] if use_bpe_for_var_name: - var_name_subtoken_maps = \ - [e.variable_name_subtoken_map for e in examples] + var_name_subtoken_maps = [e.variable_name_subtoken_map for e in examples] else: for example in examples: var_name_map = example.variable_name_map @@ -254,63 +254,73 @@ def to_batched_prediction_target(self, examples: List[Example]): else: subtoken_ids = [self.vocab.target[new_name]] var_name_subtoken_map[old_name] = subtoken_ids - var_name_subtoken_maps.append(var_name_subtoken_map) + variable_name_subtoken_maps.append(var_name_subtoken_map) max_pred_timestep = max( - sum(len(val) for val in x.values()) - for x in var_name_subtoken_maps + sum(len(val) for val in x.values()) for x in variable_name_subtoken_maps + ) + + target_variable_encoding_indices = torch.zeros( + batch_size, max_pred_timestep, dtype=torch.long + ) + target_variable_encoding_indices_mask = torch.zeros( + batch_size, max_pred_timestep + ) + + variable_tgt_name_id = torch.zeros( + batch_size, max_pred_timestep, dtype=torch.long ) - tgt_var_enc_indices = \ - torch.zeros(batch_size, max_pred_timestep, dtype=torch.long) - tgt_var_enc_indices_mask = \ - torch.zeros(batch_size, max_pred_timestep) - var_tgt_name_id = \ - torch.zeros(batch_size, max_pred_timestep, dtype=torch.long) - var_tgt_name_weight = torch.zeros(batch_size, max_pred_timestep) - var_with_new_name_mask = torch.zeros(batch_size, max_pred_timestep) + variable_tgt_name_weight = torch.zeros(batch_size, max_pred_timestep) + variable_with_new_name_mask = torch.zeros(batch_size, max_pred_timestep) auxiliary_var_mask = torch.zeros(batch_size, max_pred_timestep) - var_master_node_ptr = 0 + variable_master_node_ptr = 0 for e_id, example in enumerate(examples): ast = example.ast var_name_map = example.variable_name_map - var_ptr = 0 + + variable_ptr = 0 for var_id, var_name in enumerate(ast.variables): - new_var_name_subtoken_ids = \ - var_name_subtoken_maps[e_id][var_name] - var_end_ptr = \ - var_ptr + len(new_var_name_subtoken_ids) - var_tgt_name_id[e_id, var_ptr: var_end_ptr] = \ - torch.tensor(new_var_name_subtoken_ids, dtype=torch.long) + new_variable_name_subtoken_ids = variable_name_subtoken_maps[e_id][var_name] + variable_end_ptr = variable_ptr + len(new_var_name_subtoken_ids) + + variable_tgt_name_id[ + e_id, variable_ptr:variable_end_ptr + ] = torch.tensor(new_var_name_subtoken_ids, dtype=torch.long) + if var_name == var_name_map[var_name]: - auxiliary_var_mask[e_id, var_ptr: var_end_ptr] = 1. - var_tgt_name_weight[e_id, var_ptr: var_end_ptr] = \ - unchanged_var_weight + auxiliary_var_mask[e_id, var_ptr:var_end_ptr] = 1.0 + + variable_tgt_name_weight[ + e_id, variable_ptr:variable_end_ptr + ] = unchanged_var_weight else: - var_with_new_name_mask[e_id, var_ptr: var_end_ptr] = 1. - var_tgt_name_weight[e_id, var_ptr: var_end_ptr] = 1. + variable_with_new_name_mask[e_id, variable_ptr:variable_end_ptr] = 1.0 + variable_tgt_name_weight[e_id, variable_ptr:variable_end_ptr] = 1.0 - tgt_var_enc_indices[e_id, var_ptr: var_end_ptr] = var_id + target_variable_encoding_indices[ + e_id, variable_ptr:variable_end_ptr + ] = var_id # variable_master_node_ptr - var_master_node_ptr += 1 - var_ptr = var_end_ptr + variable_master_node_ptr += 1 + variable_ptr = variable_end_ptr - tgt_var_enc_indices_mask[e_id, :var_ptr] = 1. + target_variable_encoding_indices_mask[e_id, :variable_ptr] = 1.0 - return dict(var_tgt_name_id=var_tgt_name_id, - var_tgt_name_weight=var_tgt_name_weight, - var_with_new_name_mask=var_with_new_name_mask, - auxiliary_var_mask=auxiliary_var_mask, - target_var_encoding_indices=tgt_var_enc_indices, - target_var_encoding_indices_mask=tgt_var_enc_indices_mask) + return dict( + variable_tgt_name_id=variable_tgt_name_id, + variable_tgt_name_weight=variable_tgt_name_weight, + variable_with_new_name_mask=var_with_new_name_mask, + auxiliary_var_mask=auxiliary_var_mask, + target_variable_encoding_indices=target_variable_encoding_indices, + target_variable_encoding_indices_mask=target_variable_encoding_indices_mask, + ) -def get_json_iterator_from_tar_file(file_paths, - shuffle=False, - progress=False, - group_by=None, - buffer=True) -> Iterable: - assert group_by in (None, 'binary_file') +def get_json_iterator_from_tar_file( + file_paths, shuffle=False, progress=False, group_by=None, buffer=True +) -> Iterable: + assert group_by in (None, "binary_file") # if shuffle: # assert buffer is False @@ -323,9 +333,8 @@ def get_json_iterator_from_tar_file(file_paths, for file_path in file_paths: payloads = [] t1 = time.time() - with tarfile.open(file_path, 'r') as f: - files = [x.name for x in f.getmembers() - if x.name.endswith('.jsonl')] + with tarfile.open(file_path, "r") as f: + files = [x.name for x in f.getmembers() if x.name.endswith(".jsonl")] if progress: file_iter = tqdm(files, file=sys.stdout) else: @@ -335,43 +344,43 @@ def get_json_iterator_from_tar_file(file_paths, jsonl_file = f.extractfile(filename) if jsonl_file is not None: if group_by is None: - for line, tree_encoding_line in enumerate(jsonl_file): - payload = tree_encoding_line, \ - dict(file_name=filename, line_num=line) + for line_no, tree_encoding_line in enumerate(jsonl_file): + # if tree_encoding_line.decode().startswith('{'): + # tree_json_dict = json.loads(tree_encoding_line) + payload = ( + tree_encoding_line, + dict(file_name=filename, line_num=line_no), + ) if buffer: payloads.append(payload) else: yield payload - elif group_by == 'binary_file': - lines = [(l.decode().strip(), - dict(file_name=filename, line_num=line_no)) - for line_no, l in enumerate(jsonl_file)] + elif group_by == "binary_file": + lines = [ + ( + l.decode().strip(), + dict(file_name=filename, line_num=line_no), + ) + for line_no, l in enumerate(jsonl_file) + ] yield lines if shuffle: np.random.shuffle(payloads) - print(f'load shard {file_path} took {time.time() - t1:.4f}s', - file=sys.stderr) + print(f"load shard {file_path} took {time.time() - t1:.4f}s", file=sys.stderr) for payload in payloads: yield payload -def json_line_reader(file_path, - queue, - worker_num, - shuffle, - progress, - group_by=None, - buffer=True): - json_iterator = get_json_iterator_from_tar_file(file_path, - shuffle, - progress, - group_by=group_by, - buffer=buffer) - for json_str in json_iterator: +def json_line_reader( + file_path, queue, worker_num, shuffle, progress, group_by=None, buffer=True +): + for json_str in get_json_iterator_from_tar_file( + file_path, shuffle, progress, group_by=group_by, buffer=buffer + ): queue.put(json_str) for i in range(worker_num): @@ -379,13 +388,15 @@ def json_line_reader(file_path, def is_valid_training_example(example): - if hasattr(example, 'target_prediction_seq_length'): + if hasattr(example, "target_prediction_seq_length"): if example.target_prediction_seq_length >= 200: return False - return example.ast.size < 300 \ - and len(example.variable_name_map) > 0 \ + return ( + example.ast.size < 300 + and len(example.variable_name_map) > 0 and any(k != v for k, v in example.variable_name_map.items()) + ) def example_generator(json_queue, example_queue, consumer_num=1): @@ -396,11 +407,11 @@ def example_generator(json_queue, example_queue, consumer_num=1): json_str, meta = payload tree_json_dict = json.loads(json_str) - if 'code_tokens' in tree_json_dict: - code_tokens = tree_json_dict['code_tokens'] - example = Example.from_json_dict(tree_json_dict, - binary_file=meta, - code_tokens=code_tokens) + if "code_tokens" in tree_json_dict: + code_tokens = tree_json_dict["code_tokens"] + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, code_tokens=code_tokens + ) else: example = Example.from_json_dict(tree_json_dict, binary_file=meta) @@ -412,19 +423,21 @@ def example_generator(json_queue, example_queue, consumer_num=1): # print('[Example Generator] example generator process quit!') -def example_to_batch(json_queue, - batched_examples_queue, - batch_size, - train, - config, - worker_manager_lock, - return_examples=False, - return_prediction_target=True): +def example_to_batch( + json_queue, + batched_examples_queue, + batch_size, + train, + config, + worker_manager_lock, + return_examples=False, + return_prediction_target=True, +): batcher = Batcher(config, train) - buffer_size = config['train']['buffer_size'] + buffer_size = config["train"]["buffer_size"] buffer = [] - print(f'[ExampleToBatch] pid={os.getpid()}', file=sys.stderr) + print(f"[ExampleToBatch] pid={os.getpid()}", file=sys.stderr) def _generate_batches(): # buffer.sort(key=batcher.train_example_sort_key) @@ -434,8 +447,7 @@ def _generate_batches(): batch_examples = [] for example in buffer: - batch_size_with_example = \ - batcher.get_batch_size(batch_examples + [example]) + batch_size_with_example = batcher.get_batch_size(batch_examples + [example]) if batch_examples and batch_size_with_example > batch_size: batches.append(batch_examples) batch_examples = [] @@ -452,7 +464,7 @@ def _generate_batches(): batch = batcher.to_batch( batch_examples, return_examples=return_examples, - return_prediction_target=return_prediction_target + return_prediction_target=return_prediction_target, ) # while batched_examples_queue.qsize() > 100: # time.sleep(10) @@ -466,6 +478,8 @@ def _generate_batches(): finished = False while True: + # t1 = time.time() + while len(buffer) < buffer_size: payload = json_queue.get() if payload is None: @@ -475,14 +489,14 @@ def _generate_batches(): json_str, meta = payload tree_json_dict = json.loads(json_str) - if 'code_tokens' in tree_json_dict: - code_tokens = tree_json_dict['code_tokens'] - example = Example.from_json_dict(tree_json_dict, - binary_file=meta, - code_tokens=code_tokens) + if "code_tokens" in tree_json_dict: + code_tokens = tree_json_dict["code_tokens"] + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, code_tokens=code_tokens + ) else: - example = \ - Example.from_json_dict(tree_json_dict, binary_file=meta) + example = Example.from_json_dict(tree_json_dict, binary_file=meta) + batcher.annotate_example(example) if train: @@ -491,7 +505,9 @@ def _generate_batches(): else: buffer.append(example) + # print(f'[ExampleToBatch] {time.time() - t1}s took for loading {buffer_size} examples to buffer', file=sys.stderr) _generate_batches() + # print(f'[ExampleToBatch] {time.time() - t1}s took for batching', file=sys.stderr) if finished: break @@ -501,29 +517,28 @@ def _generate_batches(): while batcher_sync_msg.value == 0: time.sleep(1) - print(f'[ExampleToBatch] quit', file=sys.stderr) + print(f"[ExampleToBatch] quit", file=sys.stderr) sys.stderr.flush() -def worker_manager(worker_result_queue, - out_queue, - num_workers, - worker_manager_lock, - buffer_size): +def worker_manager( + worker_result_queue, out_queue, num_workers, worker_manager_lock, buffer_size +): num_finished_workers = 0 patience = 0 prev_queue_size = -1 while True: finished = False - # t0 = time.time() try: queue_size = worker_result_queue.qsize() except Exception: # just trigger data loading, for max os X queue_size = 999999 - if (queue_size > buffer_size or patience >= 10) \ - and out_queue.qsize() < buffer_size: + + if ( + queue_size > buffer_size or patience >= 10 + ) and out_queue.qsize() < buffer_size: worker_manager_lock.value = 1 patience = 0 @@ -561,7 +576,7 @@ def __init__(self, file_paths): assert isinstance(file_paths, str) self.file_paths = glob.glob(file_paths) - print(f'reading data files {self.file_paths}', file=sys.stderr) + print(f"reading data files {self.file_paths}", file=sys.stderr) example_num = 0 for _ in get_json_iterator_from_tar_file(self.file_paths): example_num += 1 @@ -573,33 +588,43 @@ def __len__(self): def __iter__(self): return self.get_iterator(progress=True) - def get_single_process_iterator(self, - shuffle=False, - progress=False) -> Iterable[Example]: - json_str_iter = \ - get_json_iterator_from_tar_file(self.file_paths, shuffle, progress) + def get_single_process_iterator( + self, shuffle=False, progress=False + ) -> Iterable[Example]: + json_str_iter = get_json_iterator_from_tar_file( + self.file_paths, shuffle, progress + ) for json_str, meta in json_str_iter: tree_json_dict = json.loads(json_str) example = Example.from_json_dict(tree_json_dict, binary_file=meta) - if example.ast.size != max(n.node_id for n in example.ast) + 1: + + if example.ast.size != max(node.node_id for node in example.ast) + 1: continue + yield example def _get_iterator(self, shuffle=False, num_workers=1): - enc_queue = multiprocessing.Queue() + json_enc_queue = multiprocessing.Queue() example_queue = multiprocessing.Queue(maxsize=5000) - args = (self.file_paths, - enc_queue, + + json_loader = multiprocessing.Process( + target=json_line_reader, + args=( + self.file_paths, + json_enc_queue, num_workers, shuffle, - False, None, False) - json_loader = \ - multiprocessing.Process(target=json_line_reader, args=args) + False, + None, + False, + ), + ) json_loader.daemon = True example_generators = [] for i in range(num_workers): - p = multiprocessing.Process(target=example_generator, - args=(enc_queue, example_queue, 1)) + p = multiprocessing.Process( + target=example_generator, args=(json_enc_queue, example_queue, 1) + ) p.daemon = True example_generators.append(p) @@ -627,15 +652,18 @@ def get_iterator(self, shuffle=False, progress=True, num_workers=1): return tqdm(iterator, total=len(self), file=sys.stdout) return iterator - def batch_iterator(self, batch_size: int, config: Dict, - return_examples=False, - return_prediction_target=None, - num_readers=3, - num_batchers=3, - progress=True, - train=False, - single_batcher=False): - # type: (...) -> Iterable[Union[Batch, Dict[str, torch.Tensor]]] + def batch_iterator( + self, + batch_size: int, + config: Dict, + return_examples=False, + return_prediction_target=None, + num_readers=3, + num_batchers=3, + progress=True, + train=False, + single_batcher=False, + ) -> Iterable[Union[Batch, Dict[str, torch.Tensor]]]: if progress: it_func = lambda x: tqdm(x, file=sys.stdout) else: @@ -643,37 +671,40 @@ def batch_iterator(self, batch_size: int, config: Dict, if single_batcher: return it_func( - self._single_process_batch_iter(batch_size, - config, - num_readers, - train) + self._single_process_batch_iter(batch_size, config, num_readers, train) ) else: return it_func( - self._batch_iterator(batch_size, - config, - num_readers, - num_batchers, - train, - return_examples, - return_prediction_target) + self._batch_iterator( + batch_size, + config, + num_readers, + num_batchers, + train, + return_examples, + return_prediction_target, + ) ) - def _batch_iterator(self, - batch_size: int, - config: Dict, - num_readers, - num_batchers, - train=False, - return_examples=False, - return_prediction_target=None) -> Iterable[Batch]: + def _batch_iterator( + self, + batch_size: int, + config: Dict, + num_readers, + num_batchers, + train=False, + return_examples=False, + return_prediction_target=None, + ) -> Iterable[Batch]: global batcher_sync_msg - batcher_sync_msg = multiprocessing.Value('i', 0) + batcher_sync_msg = multiprocessing.Value("i", 0) json_enc_queue = multiprocessing.Queue(maxsize=10000) - worker_manager_lock = multiprocessing.Value('i', 0) - args = (self.file_paths, json_enc_queue, num_readers, train, False) - json_loader = \ - multiprocessing.Process(target=json_line_reader, args=args) + worker_manager_lock = multiprocessing.Value("i", 0) + + json_loader = multiprocessing.Process( + target=json_line_reader, + args=(self.file_paths, json_enc_queue, num_readers, train, False), + ) json_loader.daemon = True example_generators = [] worker_result_queue = torch_mp.Queue(maxsize=150) @@ -682,15 +713,19 @@ def _batch_iterator(self, return_prediction_target = train for i in range(num_readers): - args = (json_enc_queue, + p = multiprocessing.Process( + target=example_to_batch, + args=( + json_enc_queue, worker_result_queue, batch_size, train, config, worker_manager_lock, return_examples, - return_prediction_target) - p = multiprocessing.Process(target=example_to_batch, args=args) + return_prediction_target, + ), + ) p.daemon = True example_generators.append(p) @@ -699,13 +734,16 @@ def _batch_iterator(self, p.start() batch_queue = queue.Queue(maxsize=100) - args = (worker_result_queue, + worker_manager_thread = threading.Thread( + target=worker_manager, + args=( + worker_result_queue, batch_queue, num_readers, worker_manager_lock, - 100) - worker_manager_thread = \ - threading.Thread(target=worker_manager, args=args) + 100, + ), + ) worker_manager_thread.start() while True: @@ -716,23 +754,25 @@ def _batch_iterator(self, yield batch worker_result_queue.close() + # start joining... batcher_sync_msg.value = 1 json_loader.join() + # json_loader quitted for p in example_generators: p.join() worker_manager_thread.join() + # example generators quitted + # batchers quiteed sys.stdout.flush() sys.stderr.flush() - def _single_process_batch_iter(self, - batch_size: int, - config: Dict, - num_readers=2, - shuffle=False) -> Iterable[Batch]: + def _single_process_batch_iter( + self, batch_size: int, config: Dict, num_readers=2, shuffle=False + ) -> Iterable[Batch]: batcher = Batcher(config) - example_iter = self.get_iterator(shuffle=shuffle, - progress=False, - num_workers=num_readers) + example_iter = self.get_iterator( + shuffle=shuffle, progress=False, num_workers=num_readers + ) # t1 = time.time() batch_examples = [] batch_node_num = 0 @@ -755,7 +795,7 @@ def _single_process_batch_iter(self, yield batch -if __name__ == '__main__': - for _example in Dataset('data/0-trees.tar.gz'): +if __name__ == "__main__": + for _example in Dataset("data/0-trees.tar.gz"): if _example.ast.size > 200: print(_example.binary_file, _example.variable_name_map) diff --git a/prediction-plugin/utils/eval_debin_prediction.py b/prediction-plugin/utils/eval_debin_prediction.py new file mode 100644 index 0000000..43842dd --- /dev/null +++ b/prediction-plugin/utils/eval_debin_prediction.py @@ -0,0 +1,95 @@ +# evaluate debin +import json + +from utils.dataset import Dataset, Example, get_json_iterator_from_tar_file +from utils.evaluation import * + + +def evaluate_debin(debin_output_path, test_example_path): + debin_predicted_data = Dataset(debin_output_path) + + debin_prediction_dict = dict() + + for example in debin_predicted_data.get_iterator(num_workers=5): + example_key = ( + example.binary_file["file_name"].split("/")[-1] + + "_" + + str(example.ast.compilation_unit) + ) + + assert example_key not in debin_prediction_dict + + pred_result = dict(example.variable_name_map.items()) + debin_prediction_dict[example_key] = pred_result + + del example + + # num_examples_without_pred = 0 + # for example in test_examples: + # example_key = example.binary_file['file_name'] + '_' + str(example.ast.compilation_unit) + # if example_key not in debin_prediction_dict: + # num_examples_without_pred += 1 + # debin_prediction_dict[example_key].variable_name_map + + # print(f'Num. examples without prediction: {num_examples_without_pred}') + + need_rename_cases = [] + func_name_in_train_acc_list = [] + func_name_not_in_train_acc_list = [] + func_body_in_train_acc_list = [] + func_body_not_in_train_acc_list = [] + debin_predicted_func_num = 0 + + test_set = Dataset(test_example_path) + + for example in test_set.get_iterator(num_workers=5): + example_key = ( + example.binary_file["file_name"].split("/")[-1] + + "_" + + str(example.ast.compilation_unit) + ) + if example_key in debin_prediction_dict: + rename_result = debin_prediction_dict[example_key] + debin_predicted_func_num += 1 + else: + rename_result = { + k: k for k in example.variable_name_map + } # use identity predictions + + example_pred_accs = [] + + for old_name, gold_new_name in example.variable_name_map.items(): + pred_new_name = rename_result[old_name] + var_metric = Evaluator.get_soft_metrics(pred_new_name, gold_new_name) + + example_pred_accs.append(var_metric) + + if gold_new_name != old_name: # and gold_new_name in model.vocab.target: + need_rename_cases.append(var_metric) + + if example.test_meta["function_name_in_train"]: + func_name_in_train_acc_list.append(var_metric) + else: + func_name_not_in_train_acc_list.append(var_metric) + + if example.test_meta["function_body_in_train"]: + func_body_in_train_acc_list.append(var_metric) + else: + func_body_not_in_train_acc_list.append(var_metric) + + eval_results = dict( + corpus_need_rename_acc=Evaluator.average(need_rename_cases), + func_name_in_train_acc=Evaluator.average(func_name_in_train_acc_list), + func_name_not_in_train_acc=Evaluator.average(func_name_not_in_train_acc_list), + func_body_in_train_acc=Evaluator.average(func_body_in_train_acc_list), + func_body_not_in_train_acc=Evaluator.average(func_body_not_in_train_acc_list), + ) + + print("Num. functions debin predicted: ", debin_predicted_func_num) + print(eval_results) + + +if __name__ == "__main__": + evaluate_debin( + "data/debin.predictions.0.01.tar", "data/all_trees_tokenized_0410/test.tar" + ) diff --git a/prediction-plugin/utils/evaluation.py b/prediction-plugin/utils/evaluation.py index 605fabf..42827b2 100644 --- a/prediction-plugin/utils/evaluation.py +++ b/prediction-plugin/utils/evaluation.py @@ -1,13 +1,11 @@ -from typing import Dict, List, Any +from typing import Any, Dict, List, Callable +import editdistance import numpy as np import torch - -import editdistance - +from model.model import RenamingModel from utils import nn_util from utils.dataset import Dataset -from model.model import RenamingModel class Evaluator(object): @@ -17,10 +15,9 @@ def get_soft_metrics(pred_name: str, gold_name: str) -> Dict: cer = float(edit_distance / len(gold_name)) acc = float(pred_name == gold_name) - return dict(edit_distance=edit_distance, - ref_len=len(gold_name), - cer=cer, - accuracy=acc) + return dict( + edit_distance=edit_distance, ref_len=len(gold_name), cer=cer, accuracy=acc + ) @staticmethod def average(metrics_list: List[Dict]) -> Dict: @@ -30,7 +27,7 @@ def average(metrics_list: List[Dict]) -> Dict: agg_results.setdefault(key, []).append(val) avg_results = dict() - avg_results['corpus_cer'] = 0 + avg_results["corpus_cer"] = 0 for key, val in agg_results.items(): avg_results[key] = np.average(val) @@ -38,36 +35,39 @@ def average(metrics_list: List[Dict]) -> Dict: return avg_results @staticmethod - def evaluate_ppl(model: RenamingModel, - dataset: Dataset, - config: Dict, - predicate: Any = None): - if predicate is None: - def predicate(_): - return True - - eval_batch_size = config['train']['batch_size'] - num_readers = config['train']['num_readers'] - num_batchers = config['train']['num_batchers'] - data_iter = dataset.batch_iterator(batch_size=eval_batch_size, - train=False, progress=True, - return_examples=False, - return_prediction_target=True, - config=model.config, - num_readers=num_readers, - num_batchers=num_batchers) + def _always_true(_: Any) -> bool: + return True + + @staticmethod + def evaluate_ppl( + model: RenamingModel, dataset: Dataset, config: Dict, predicate: Callable[[Any], bool] = _always_true + ): + eval_batch_size = config["train"]["batch_size"] + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=False, + return_prediction_target=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) was_training = model.training model.eval() - cum_log_probs = 0. + cum_log_probs = 0.0 cum_num_examples = 0 with torch.no_grad(): for batch in data_iter: - td = batch.tensor_dict - nn_util.to(td, model.device) - result = model(td, td['prediction_target']) - log_probs = result['batch_log_prob'].cpu().tolist() - for e_id, test_meta in enumerate(td['test_meta']): + nn_util.to(batch.tensor_dict, model.device) + result = model( + batch.tensor_dict, batch.tensor_dict["prediction_target"] + ) + log_probs = result["batch_log_prob"].cpu().tolist() + for e_id, test_meta in enumerate(batch.tensor_dict["test_meta"]): if predicate(test_meta): log_prob = log_probs[e_id] cum_log_probs += log_prob @@ -81,24 +81,25 @@ def predicate(_): return ppl @staticmethod - def decode(model: RenamingModel, - dataset: Dataset, - config: Dict, - eval_batch_size=None): + def decode( + model: RenamingModel, dataset: Dataset, config: Dict, eval_batch_size=None + ): if eval_batch_size is None: - if 'eval_batch_size' in config['train']: - eval_batch_size = config['train']['eval_batch_size'] + if "eval_batch_size" in config["train"]: + eval_batch_size = config["train"]["eval_batch_size"] else: - eval_batch_size = config['train']['batch_size'] - num_readers = config['train']['num_readers'] - num_batchers = config['train']['num_batchers'] - data_iter = dataset.batch_iterator(batch_size=eval_batch_size, - train=False, - progress=True, - return_examples=True, - config=model.config, - num_readers=num_readers, - num_batchers=num_batchers) + eval_batch_size = config["train"]["batch_size"] + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) model.eval() all_examples = dict() @@ -109,40 +110,48 @@ def decode(model: RenamingModel, for example, rename_result in zip(examples, rename_results): example_pred_accs = [] top_rename_result = rename_result[0] - for old_name, gold_new_name \ - in example.variable_name_map.items(): + for old_name, gold_new_name in example.variable_name_map.items(): pred = top_rename_result[old_name] - pred_new_name = pred['new_name'] - var_metric = Evaluator.get_soft_metrics(pred_new_name, - gold_new_name) + pred_new_name = pred["new_name"] + var_metric = Evaluator.get_soft_metrics( + pred_new_name, gold_new_name + ) example_pred_accs.append(var_metric) - file_name = example.binary_file['file_name'] - line_num = example.binary_file['line_num'] + file_name = example.binary_file["file_name"] + line_num = example.binary_file["line_num"] fun_name = example.ast.compilation_unit - all_examples[f'{file_name}_{line_num}_{fun_name}'] = \ - (rename_result, Evaluator.average(example_pred_accs)) + all_examples[f"{file_name}_{line_num}_{fun_name}"] = ( + rename_result, + Evaluator.average(example_pred_accs), + ) return all_examples @staticmethod - def decode_and_evaluate(model: RenamingModel, - dataset: Dataset, - config: Dict, - return_results=False, - eval_batch_size=None): + def decode_and_evaluate( + model: RenamingModel, + dataset: Dataset, + config: Dict, + return_results=False, + eval_batch_size=None, + ): if eval_batch_size is None: - if 'eval_batch_size' in config['train']: - eval_batch_size = config['train']['eval_batch_size'] + if "eval_batch_size" in config["train"]: + eval_batch_size = config["train"]["eval_batch_size"] else: - eval_batch_size = config['train']['batch_size'] - num_readers = config['train']['num_readers'] - num_batchers = config['train']['num_batchers'] - data_iter = dataset.batch_iterator(batch_size=eval_batch_size, - train=False, progress=True, - return_examples=True, - config=model.config, - num_readers=num_readers, - num_batchers=num_batchers) + eval_batch_size = config["train"]["batch_size"] + + num_readers = config["train"]["num_readers"] + num_batchers = config["train"]["num_batchers"] + data_iter = dataset.batch_iterator( + batch_size=eval_batch_size, + train=False, + progress=True, + return_examples=True, + config=model.config, + num_readers=num_readers, + num_batchers=num_batchers, + ) was_training = model.training model.eval() @@ -165,24 +174,24 @@ def decode_and_evaluate(model: RenamingModel, example_pred_accs = [] top_rename_result = rename_result[0] - for old_name, gold_new_name \ - in example.variable_name_map.items(): + for old_name, gold_new_name in example.variable_name_map.items(): pred = top_rename_result[old_name] - pred_new_name = pred['new_name'] - var_metric = Evaluator.get_soft_metrics(pred_new_name, - gold_new_name) + pred_new_name = pred["new_name"] + var_metric = Evaluator.get_soft_metrics( + pred_new_name, gold_new_name + ) # is_correct = pred_new_name == gold_new_name example_pred_accs.append(var_metric) if gold_new_name != old_name: need_rename_cases.append(var_metric) - if example.test_meta['function_name_in_train']: + if example.test_meta["function_name_in_train"]: func_name_in_train_acc.append(var_metric) else: func_name_not_in_train_acc.append(var_metric) - if example.test_meta['function_body_in_train']: + if example.test_meta["function_body_in_train"]: func_body_in_train_acc.append(var_metric) else: func_body_not_in_train_acc.append(var_metric) @@ -191,12 +200,15 @@ def decode_and_evaluate(model: RenamingModel, example_acc_list.append(example_pred_accs) if return_results: - example = \ - f"{example.binary_file['file_name']}_" \ + example_key = ( + f"{example.binary_file['file_name']}_" f"{example.binary_file['line_num']}" - all_examples[example] = \ - (rename_result, - Evaluator.average(example_pred_accs)) + ) + + all_examples[example_key] = ( + rename_result, + Evaluator.average(example_pred_accs), + ) valid_example_num = len(example_acc_list) num_variables = len(variable_acc_list) @@ -210,19 +222,21 @@ def decode_and_evaluate(model: RenamingModel, name_not_in_train_acc = Evaluator.average(func_name_not_in_train_acc) body_in_train_acc = Evaluator.average(func_body_in_train_acc) body_not_in_train_acc = Evaluator.average(func_body_not_in_train_acc) - eval_results = dict(corpus_acc=corpus_acc, - corpus_need_rename_acc=need_rename_acc, - func_name_in_train_acc=name_in_train_acc, - func_name_not_in_train_acc=name_not_in_train_acc, - func_body_in_train_acc=body_in_train_acc, - func_body_not_in_train_acc=body_not_in_train_acc, - num_variables=num_variables, - num_valid_examples=valid_example_num) + eval_results = dict( + corpus_acc=corpus_acc, + corpus_need_rename_acc=need_rename_acc, + func_name_in_train_acc=name_in_train_acc, + func_name_not_in_train_acc=name_not_in_train_acc, + func_body_in_train_acc=body_in_train_acc, + func_body_not_in_train_acc=body_not_in_train_acc, + num_variables=num_variables, + num_valid_examples=valid_example_num, + ) if return_results: return eval_results, all_examples return eval_results -if __name__ == '__main__': - print(Evaluator.get_soft_metrics('file_name', 'filename')) +if __name__ == "__main__": + print(Evaluator.get_soft_metrics("file_name", "filename")) diff --git a/prediction-plugin/utils/get_stat.py b/prediction-plugin/utils/get_stat.py new file mode 100644 index 0000000..f403279 --- /dev/null +++ b/prediction-plugin/utils/get_stat.py @@ -0,0 +1,59 @@ +# compute training data statistics +import gc +from collections import Counter + +import numpy as np +from utils.dataset import Dataset + +gc.collect() + + +def compute_dataset_stat(): + train_path = "data/all_trees_tokenized_0410/train-shard-*.tar" + train_set = Dataset(train_path) + + num_train_examples = 0 + num_variables = [] + num_variables_with_new_name = [] + ast_sizes = [] + var_name_freq = Counter() + + dev_set = Dataset("data/all_trees_tokenized_0410/dev.tar") + test_set = Dataset("data/all_trees_tokenized_0410/test.tar") + + for dataset_id, dataset in enumerate([train_set, dev_set, test_set]): + for example in dataset.get_iterator(num_workers=5): + example_num_variables = len(example.variable_name_map) + + n_var_new_name = 0 + for old_var_name, new_var_name in example.variable_name_map.items(): + if old_var_name != new_var_name: + n_var_new_name += 1 + + if dataset_id == 0: + var_name_freq[new_var_name] += 1 + + num_variables.append(example_num_variables) + num_variables_with_new_name.append(n_var_new_name) + ast_sizes.append(example.ast.size) + + del example + + with open("data/all_trees_tokenized_0410/train_var_name_freq.txt", "w") as f: + for var_name, freq in var_name_freq.most_common(): + f.write(f"{var_name}\t{freq}\n") + + with open("data/all_trees_tokenized_0410/stat.txt", "w") as f: + print(len(train_set), len(dev_set), len(test_set), file=f) + print("num total functions: ", num_train_examples, file=f) + print("avg. size of AST:", np.average(ast_sizes), file=f) + print("avg. num of variables:", np.average(num_variables), file=f) + print( + "avg. num of variables with new name:", + np.average(num_variables_with_new_name), + file=f, + ) + + +if __name__ == "__main__": + compute_dataset_stat() diff --git a/prediction-plugin/utils/grammar.py b/prediction-plugin/utils/grammar.py index aeb5fbc..54ec9f3 100644 --- a/prediction-plugin/utils/grammar.py +++ b/prediction-plugin/utils/grammar.py @@ -1,18 +1,18 @@ class Grammar(object): def __init__(self, syntax_types, variable_types): self.syntax_types = list(sorted(syntax_types)) - self.variable_types = ['', ''] + list(sorted(variable_types)) + self.variable_types = ["", ""] + list(sorted(variable_types)) - self.syntax_type_to_id = \ - {type: id for id, type in enumerate(self.syntax_types)} - self._variable_type_to_id = \ - {type: id for id, type in enumerate(self.variable_types)} + self.syntax_type_to_id = {type: id for id, type in enumerate(self.syntax_types)} + self._variable_type_to_id = { + type: id for id, type in enumerate(self.variable_types) + } def variable_type_to_id(self, type_token): if type_token in self._variable_type_to_id: return self._variable_type_to_id[type_token] - return self._variable_type_to_id[''] + return self._variable_type_to_id[""] @property def params(self): @@ -20,6 +20,6 @@ def params(self): @classmethod def load(cls, params): - grammar = cls(params['syntax_types'], params['variable_types']) + grammar = cls(params["syntax_types"], params["variable_types"]) return grammar diff --git a/prediction-plugin/utils/graph.py b/prediction-plugin/utils/graph.py index dcff86f..23f3cae 100644 --- a/prediction-plugin/utils/graph.py +++ b/prediction-plugin/utils/graph.py @@ -1,5 +1,5 @@ -from collections import defaultdict, OrderedDict -from typing import List, Dict +from collections import OrderedDict, defaultdict +from typing import Dict, List from utils.ast import AbstractSyntaxTree, SyntaxNode @@ -24,11 +24,9 @@ def register_tree(self, ast: AbstractSyntaxTree): for node in ast: self.register_node(ast_id, node.node_id) - def register_node(self, - tree_id, - node, - group='ast_nodes', - return_node_index_in_group=False): + def register_node( + self, tree_id, node, group="ast_nodes", return_node_index_in_group=False + ): if group not in self.node_groups[tree_id]: self.node_groups[tree_id][group] = OrderedDict() @@ -44,7 +42,7 @@ def register_node(self, return packed_node_id - def get_packed_node_id(self, tree_id, node, group='ast_nodes'): + def get_packed_node_id(self, tree_id, node, group="ast_nodes"): if isinstance(node, SyntaxNode): node = node.node_id return self.node_groups[tree_id][group][node] diff --git a/prediction-plugin/utils/gz_to_jsonl.py b/prediction-plugin/utils/gz_to_jsonl.py new file mode 100644 index 0000000..0858e97 --- /dev/null +++ b/prediction-plugin/utils/gz_to_jsonl.py @@ -0,0 +1,17 @@ +import json +import sys +import tarfile + +from tqdm import tqdm + +file_name = sys.argv[1] +tgt_file_name = sys.argv[2] + +with open(tgt_file_name, "w") as f_tgt: + with tarfile.open(file_name, "r") as f: + for filename in tqdm(f.getmembers()): + json_file = f.extractfile(filename) + if json_file is not None: + json_str = json_file.read() + json_str = json.dumps(json.loads(json_str)) + f_tgt.write(json_str + "\n") diff --git a/prediction-plugin/utils/lexer.py b/prediction-plugin/utils/lexer.py index 2cd1106..4ca4d4b 100644 --- a/prediction-plugin/utils/lexer.py +++ b/prediction-plugin/utils/lexer.py @@ -7,10 +7,10 @@ # reference shadowed global variables. We throw away this operator. from enum import Enum, auto + from pygments import lex -from pygments.token import Token -from pygments.token import is_token_subtype from pygments.lexers.c_cpp import CLexer, inherit +from pygments.token import Token, is_token_subtype Token.Placeholder = Token.Token.Placeholder @@ -42,8 +42,7 @@ def get_tokens(self, var_names=Names.RAW): # Pygments breaks up strings into individual tokens representing # things like opening quotes and escaped characters. We want to # collapse all of these into a single string literal token. - if previous_string \ - and not is_token_subtype(token_type, Token.String): + if previous_string and not is_token_subtype(token_type, Token.String): yield (Token.String, previous_string) previous_string = None if is_token_subtype(token_type, Token.String): @@ -57,15 +56,14 @@ def get_tokens(self, var_names=Names.RAW): elif is_token_subtype(token_type, Token.Comment): continue # Skip the :: token added by HexRays - elif is_token_subtype(token_type, Token.Operator) \ - and token == '::': + elif is_token_subtype(token_type, Token.Operator) and token == "::": continue # Replace the text of placeholder tokens elif is_token_subtype(token_type, Token.Placeholder): yield { Names.RAW: (token_type, token), - Names.SOURCE: (token_type, token.split('@@')[2]), - Names.TARGET: (token_type, token.split('@@')[3]), + Names.SOURCE: (token_type, token.split("@@")[2]), + Names.TARGET: (token_type, token.split("@@")[3]), }[var_names] elif not is_token_subtype(token_type, Token.Text): yield (token_type, token.strip()) @@ -79,39 +77,39 @@ def get_tokens(self, var_names=Names.RAW): class HexRaysLexer(CLexer): # Additional tokens tokens = { - 'statements': [ - (r'->', Token.Operator), - (r'\+\+', Token.Operator), - (r'--', Token.Operator), - (r'==', Token.Operator), - (r'!=', Token.Operator), - (r'>=', Token.Operator), - (r'<=', Token.Operator), - (r'&&', Token.Operator), - (r'\|\|', Token.Operator), - (r'\+=', Token.Operator), - (r'-=', Token.Operator), - (r'\*=', Token.Operator), - (r'/=', Token.Operator), - (r'%=', Token.Operator), - (r'&=', Token.Operator), - (r'\^=', Token.Operator), - (r'\|=', Token.Operator), - (r'<<=', Token.Operator), - (r'>>=', Token.Operator), - (r'<<', Token.Operator), - (r'>>', Token.Operator), - (r'\.\.\.', Token.Operator), - (r'##', Token.Operator), - (r'::', Token.Operator), - (r'@@VAR_[0-9]+@@\w+@@\w+', Token.Placeholder.Var), - inherit + "statements": [ + (r"->", Token.Operator), + (r"\+\+", Token.Operator), + (r"--", Token.Operator), + (r"==", Token.Operator), + (r"!=", Token.Operator), + (r">=", Token.Operator), + (r"<=", Token.Operator), + (r"&&", Token.Operator), + (r"\|\|", Token.Operator), + (r"\+=", Token.Operator), + (r"-=", Token.Operator), + (r"\*=", Token.Operator), + (r"/=", Token.Operator), + (r"%=", Token.Operator), + (r"&=", Token.Operator), + (r"\^=", Token.Operator), + (r"\|=", Token.Operator), + (r"<<=", Token.Operator), + (r">>=", Token.Operator), + (r"<<", Token.Operator), + (r">>", Token.Operator), + (r"\.\.\.", Token.Operator), + (r"##", Token.Operator), + (r"::", Token.Operator), + (r"@@VAR_[0-9]+@@\w+@@\w+", Token.Placeholder.Var), + inherit, ] } -if __name__ == '__main__': - code = ''' +if __name__ == "__main__": + code = """ __int64 (__fastcall **)(unsigned __int64, signed __int64, __int64, _QWORD, __int64, signed __int64) { int a = "asdfsdf"; @@ -119,5 +117,5 @@ class HexRaysLexer(CLexer): a = asd::safd(); sadf=12 /*asdf*/ -}''' +}""" print([str(token[0]) for token in Lexer(code).get_tokens()]) diff --git a/prediction-plugin/utils/nn_util.py b/prediction-plugin/utils/nn_util.py index 46cbafe..630a729 100644 --- a/prediction-plugin/utils/nn_util.py +++ b/prediction-plugin/utils/nn_util.py @@ -4,7 +4,6 @@ import numpy as np import torch - SMALL_NUMBER = 1e-8 @@ -15,8 +14,8 @@ def glorot_init(params): def to(data, device: torch.device): - if 'adj_lists' in data: - [x.to(device) for x in data['adj_lists']] + if "adj_lists" in data: + [x.to(device) for x in data["adj_lists"]] if isinstance(data, dict): for key, val in data.items(): @@ -44,7 +43,7 @@ def batch_iter(data, batch_size, shuffle=False): np.random.shuffle(index_array) for i in range(batch_num): - indices = index_array[i * batch_size: (i + 1) * batch_size] + indices = index_array[i * batch_size : (i + 1) * batch_size] examples = [data[idx] for idx in indices] yield examples @@ -65,24 +64,23 @@ def get_tensor_dict_size(tensor_dict): return total_num_elements -def dot_prod_attention(h_t: torch.Tensor, - src_encoding: torch.Tensor, - src_encoding_att_linear: torch.Tensor, - mask: torch.Tensor = None): - # type: (...) -> Tuple[torch.Tensor, torch.Tensor] - att_weight = \ - torch.bmm(src_encoding_att_linear, h_t.unsqueeze(2)).squeeze(2) +def dot_prod_attention( + h_t: torch.Tensor, + src_encoding: torch.Tensor, + src_encoding_att_linear: torch.Tensor, + mask: torch.Tensor = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + # (batch_size, src_sent_len) + att_weight = torch.bmm(src_encoding_att_linear, h_t.unsqueeze(2)).squeeze(2) + if mask is not None: - att_weight.data.masked_fill_((1. - mask).bool(), -float('inf')) + att_weight.data.masked_fill_((1.0 - mask).bool(), -float("inf")) softmaxed_att_weight = torch.softmax(att_weight, dim=-1) att_view = (att_weight.size(0), 1, att_weight.size(1)) # (batch_size, hidden_size) - ctx_vec = torch.bmm( - softmaxed_att_weight.view(*att_view), - src_encoding - ).squeeze(1) + ctx_vec = torch.bmm(softmaxed_att_weight.view(*att_view), src_encoding).squeeze(1) return ctx_vec, softmaxed_att_weight @@ -105,7 +103,8 @@ def get_lengths_from_binary_sequence_mask(mask: torch.Tensor): def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor): - """Sort a batch first tensor by some specified lengths. + """ + Sort a batch first tensor by some specified lengths. Parameters ---------- tensor : torch.FloatTensor, required. @@ -122,32 +121,31 @@ def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor): The original sequence_lengths sorted by decreasing size. restoration_indices : torch.LongTensor Indices into the sorted_tensor such that - ``sorted_tensor.index_select(0, restoration_indices) == - original_tensor`` + ``sorted_tensor.index_select(0, restoration_indices) == original_tensor`` permuation_index : torch.LongTensor The indices used to sort the tensor. This is useful if you want to sort many tensors using the same ordering. """ - if not isinstance(tensor, torch.Tensor) \ - or not isinstance(sequence_lengths, torch.Tensor): - raise ValueError( - "Both the tensor and sequence lengths must be torch.Tensors." - ) + if not isinstance(tensor, torch.Tensor) or not isinstance( + sequence_lengths, torch.Tensor + ): + raise ValueError("Both the tensor and sequence lengths must be torch.Tensors.") - sorted_sequence_lengths, permutation_index = \ - sequence_lengths.sort(0, descending=True) + sorted_sequence_lengths, permutation_index = sequence_lengths.sort( + 0, descending=True + ) sorted_tensor = tensor.index_select(0, permutation_index) - index_range = torch.arange( - 0, len(sequence_lengths), device=sequence_lengths.device - ) + index_range = torch.arange(0, len(sequence_lengths), device=sequence_lengths.device) # This is the equivalent of zipping with index, sorting by the original # sequence lengths and returning the now sorted indices. _, reverse_mapping = permutation_index.sort(0, descending=False) restoration_indices = index_range.index_select(0, reverse_mapping) - return (sorted_tensor, - sorted_sequence_lengths, - restoration_indices, - permutation_index) + return ( + sorted_tensor, + sorted_sequence_lengths, + restoration_indices, + permutation_index, + ) diff --git a/prediction-plugin/utils/preprocess.py b/prediction-plugin/utils/preprocess.py index 29ca03a..569f3a1 100644 --- a/prediction-plugin/utils/preprocess.py +++ b/prediction-plugin/utils/preprocess.py @@ -5,18 +5,33 @@ Options: -h --help Show this screen. + --shard-size= shard size [default: 3000] + --test-file= test file + --no-filtering do not filter files """ import glob import multiprocessing +import os +import random +import sys +import tarfile +from collections import Iterable +from multiprocessing import Process +from typing import Tuple + +import numpy as np import ujson as json - from docopt import docopt -import os - +from tqdm import tqdm from utils.ast import SyntaxNode -from utils.code_processing import \ - canonicalize_code, preprocess_ast, tokenize_raw_code +from utils.code_processing import ( + annotate_type, + canonicalize_code, + canonicalize_constants, + preprocess_ast, + tokenize_raw_code, +) from utils.dataset import Example, json_line_reader all_functions = dict() # indexed by binaries @@ -24,34 +39,36 @@ def is_valid_example(example): try: - is_valid = example.ast.size < 300 and \ - len(example.variable_name_map) > 0 + is_valid = example.ast.size < 300 and len(example.variable_name_map) > 0 except RecursionError: is_valid = False return is_valid + def generate_example(json_str, binary_file): tree_json_dict = json.loads(json_str) - root = SyntaxNode.from_json_dict(tree_json_dict['ast']) + root = SyntaxNode.from_json_dict(tree_json_dict["ast"]) - preprocess_ast(root, code=tree_json_dict['raw_code']) - code_tokens = tokenize_raw_code(tree_json_dict['raw_code']) - tree_json_dict['code_tokens'] = code_tokens + preprocess_ast(root, code=tree_json_dict["raw_code"]) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) + tree_json_dict["code_tokens"] = code_tokens # add function name to the name field of the root block - root.name = tree_json_dict['function'] - root.named_fields.add('name') + root.name = tree_json_dict["function"] + root.named_fields.add("name") new_json_dict = root.to_json_dict() - tree_json_dict['ast'] = new_json_dict + tree_json_dict["ast"] = new_json_dict json_str = json.dumps(tree_json_dict) - example = Example.from_json_dict(tree_json_dict, - binary_file=binary_file, - json_str=json_str, - code_tokens=code_tokens) + example = Example.from_json_dict( + tree_json_dict, + binary_file=binary_file, + json_str=json_str, + code_tokens=code_tokens, + ) if True or is_valid_example(example): canonical_code = canonicalize_code(example.ast.code) @@ -60,7 +77,9 @@ def generate_example(json_str, binary_file): else: return None + def example_generator(json_queue, example_queue, args, consumer_num=1): + enable_filter = not args["--no-filtering"] while True: payload = json_queue.get() if payload is None: @@ -70,27 +89,28 @@ def example_generator(json_queue, example_queue, args, consumer_num=1): for json_str, meta in payload: tree_json_dict = json.loads(json_str) - root = SyntaxNode.from_json_dict(tree_json_dict['ast']) + root = SyntaxNode.from_json_dict(tree_json_dict["ast"]) # root_reconstr = SyntaxNode.from_json_dict(root.to_json_dict()) # assert root == root_reconstr - preprocess_ast(root, code=tree_json_dict['raw_code']) - code_tokens = tokenize_raw_code(tree_json_dict['raw_code']) - tree_json_dict['code_tokens'] = code_tokens + preprocess_ast(root, code=tree_json_dict["raw_code"]) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) + tree_json_dict["code_tokens"] = code_tokens # add function name to the name field of the root block - root.name = tree_json_dict['function'] - root.named_fields.add('name') + root.name = tree_json_dict["function"] + root.named_fields.add("name") new_json_dict = root.to_json_dict() - tree_json_dict['ast'] = new_json_dict + tree_json_dict["ast"] = new_json_dict json_str = json.dumps(tree_json_dict) - example = Example.from_json_dict(tree_json_dict, - binary_file=meta, - json_str=json_str) + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, json_str=json_str + ) - if is_valid_example(example): + is_valid = is_valid_example(example) if enable_filter else True + if is_valid: canonical_code = canonicalize_code(example.ast.code) example.canonical_code = canonical_code examples.append(example) @@ -100,23 +120,27 @@ def example_generator(json_queue, example_queue, args, consumer_num=1): for i in range(consumer_num): example_queue.put(None) - print('example generator quit!') + print("example generator quit!") def main(args): - tgt_folder = args['TARGET_FOLDER'] - pattern_list = args['TAR_FILES'].split(',') + test_run = bool(args["--test-file"]) + if test_run: + np.random.seed(1234) + random.seed(1992) + + tgt_folder = args["TARGET_FOLDER"] + pattern_list = args["TAR_FILES"].split(",") tar_files = [] for pattern in pattern_list: tar_files.extend(glob.glob(pattern)) print(tar_files) - os.system(f'mkdir -p {tgt_folder}') - os.system(f'mkdir -p {tgt_folder}/files') + os.system(f'mkdir -p "{tgt_folder}/files"') num_workers = 14 for tar_file in tar_files: - print(f'read {tar_file}') + print(f"read {tar_file}") valid_example_count = 0 json_enc_queue = multiprocessing.Queue() @@ -124,12 +148,14 @@ def main(args): json_loader = multiprocessing.Process( target=json_line_reader, - args=(os.path.expanduser(tar_file), - json_enc_queue, - num_workers, - False, - False, - 'binary_file') + args=( + os.path.expanduser(tar_file), + json_enc_queue, + num_workers, + False, + False, + "binary_file", + ), ) json_loader.daemon = True json_loader.start() @@ -137,8 +163,7 @@ def main(args): example_generators = [] for i in range(num_workers): p = multiprocessing.Process( - target=example_generator, - args=(json_enc_queue, example_queue, args, 1) + target=example_generator, args=(json_enc_queue, example_queue, args, 1) ) p.daemon = True p.start() @@ -148,7 +173,7 @@ def main(args): while True: payload = example_queue.get() if payload is None: - print('received None!') + print("received None!") n_finished += 1 if n_finished == num_workers: break @@ -157,21 +182,17 @@ def main(args): examples = payload if examples: - json_file_name = \ - examples[0].binary_file['file_name'].split('/')[-1] - with open(os.path.join(tgt_folder, 'files/', json_file_name), - 'w') as f: + json_file_name = examples[0].binary_file["file_name"].split("/")[-1] + with open(os.path.join(tgt_folder, "files/", json_file_name), "w") as f: for example in examples: - f.write(example.json_str + '\n') - all_functions.setdefault( - json_file_name, - dict() - )[example.ast.compilation_unit] = \ - example.canonical_code + f.write(example.json_str + "\n") + all_functions.setdefault(json_file_name, dict())[ + example.ast.compilation_unit + ] = example.canonical_code valid_example_count += len(examples) - print('valid examples: ', valid_example_count) + print("valid examples: ", valid_example_count) json_enc_queue.close() example_queue.close() @@ -181,17 +202,17 @@ def main(args): p.join() cur_dir = os.getcwd() - all_files = glob.glob(os.path.join(tgt_folder, 'files/*.jsonl')) + all_files = glob.glob(os.path.join(tgt_folder, "files/*.jsonl")) sorted(all_files) # sort all files by names all_files = list(all_files) file_num = len(all_files) - print(f'{file_num} valid binary files.') + print(f"{file_num} valid binary files.") - def _dump_file(tgt_file_name, file_names): - with open(os.path.join(tgt_folder, 'file_list.txt'), 'w') as f: + def _dump_file(tgt_file_name, file_names, dev=False): + with open(os.path.join(tgt_folder, "file_list.txt"), "w") as f: for file_name in file_names: - last_file_name = file_name.split('/')[-1] - f.write(last_file_name + '\n') + last_file_name = file_name.split("/")[-1] + f.write(last_file_name + "\n") with open(file_name) as fr: all_lines = fr.readlines() @@ -199,22 +220,44 @@ def _dump_file(tgt_file_name, file_names): replace_lines = [] for line in all_lines: json_dict = json.loads(line.strip()) + if dev: + func_name = json_dict["function"] + canonical_code = all_functions[last_file_name][func_name] + func_name_in_train = False + func_body_in_train = False + if func_name in train_functions: + func_name_in_train = True + if canonical_code in train_functions[func_name]: + func_body_in_train = True + + json_dict["test_meta"] = dict( + function_name_in_train=func_name_in_train, + function_body_in_train=func_body_in_train, + ) new_json_str = json.dumps(json_dict) replace_lines.append(new_json_str.strip()) - with open(file_name, 'w') as fw: + with open(file_name, "w") as fw: for line in replace_lines: - fw.write(line + '\n') + fw.write(line + "\n") - os.chdir(os.path.join(tgt_folder, 'files')) - print('creating tar file...') - os.system(f'tar cf ../{tgt_file_name} -T ../file_list.txt') + os.chdir(os.path.join(tgt_folder, "files")) + print("creating tar file...") + os.system(f"tar cf ../{tgt_file_name} -T ../file_list.txt") os.chdir(cur_dir) - print('dumping files') - _dump_file('preprocessed.tar', all_files) + if test_run: + print("dumping dev files") + _dump_dev_file("dev.tar", dev_files, dev=True) + + print("dumping test files") + _dump_dev_file("test.tar", test_files, dev=True) + else: + print("dumping files") + _dump_file("preprocessed.tar", all_files) + -if __name__ == '__main__': +if __name__ == "__main__": args = docopt(__doc__) main(args) diff --git a/prediction-plugin/utils/sequential_preprocess.py b/prediction-plugin/utils/sequential_preprocess.py new file mode 100644 index 0000000..c19748a --- /dev/null +++ b/prediction-plugin/utils/sequential_preprocess.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python +""" +Usage: + preprocess.py [options] TAR_FILES OUTPUT_CODE_FILE + +Options: + -h --help Show this screen. +""" + +import ujson as json +from docopt import docopt +from utils.code_processing import VARIABLE_ANNOTATION +from utils.dataset import Dataset +from utils.lexer import * + + +def processor(input_queue, output_queue): + while True: + payload = input_queue.get() + if payload is None: + break + + examples = [] + for json_str, meta in payload: + tree_json_dict = json.loads(json_str) + code_tokens = tokenize_raw_code(tree_json_dict["raw_code"]) + + tree_json_dict["code_tokens"] = code_tokens + json_str = json.dumps(tree_json_dict) + + preprocess_ast(root, code=tree_json_dict["raw_code"]) + + # add function name to the name field of the root block + root.name = tree_json_dict["function"] + root.named_fields.add("name") + + new_json_dict = root.to_json_dict() + tree_json_dict["ast"] = new_json_dict + + example = Example.from_json_dict( + tree_json_dict, binary_file=meta, json_str=json_str + ) + + if is_valid_example(example): + canonical_code = canonicalize_code(example.ast.code) + example.canonical_code = canonical_code + examples.append(example) + + example_queue.put(examples) + + for i in range(consumer_num): + example_queue.put(None) + + print("example generator quited!") + + +def main(args): + dataset = Dataset(args["TAR_FILES"]) + code_line_file = open(args["OUTPUT_CODE_FILE"], "w") + all_preserved_tokens = set() + for example in dataset.get_iterator(num_workers=5): + code = example.ast.code + # code_tokens = tokenize_raw_code(code) + # preserved_tokens = [token for token in code_tokens if token.startswith('@@') and token.endswith('@@')] + # all_preserved_tokens.update(preserved_tokens) + + # code_line_file.write(' '.join(code_tokens) + '\n') + + code_line_file.close() + + with open(args["OUTPUT_CODE_FILE"] + ".preserved_tokens.txt", "w") as f: + for token in all_preserved_tokens: + f.write(token + "\n") + + +if __name__ == "__main__": + args = docopt(__doc__) + main(args) diff --git a/prediction-plugin/utils/subsample.py b/prediction-plugin/utils/subsample.py new file mode 100644 index 0000000..df672da --- /dev/null +++ b/prediction-plugin/utils/subsample.py @@ -0,0 +1,51 @@ +import glob +import os +import sys + +import numpy as np +from sh import tar + + +def main(): + tar_files = sys.argv[1] + sample_ratio = float(sys.argv[2]) + + np.random.seed(1234) + + tar_files = glob.glob(tar_files) + print(tar_files) + + work_dir = os.path.dirname(tar_files[0]) + os.system(f"mkdir -p {work_dir}/binary_files/") + + for tar_file in tar_files: + print(f"extracting {tar_file}") + tar("-C", f"{work_dir}/binary_files/", "-xvf", tar_file) + + all_files = glob.glob(f"{work_dir}/binary_files/*.jsonl") + all_files.sort() + + print(f"{len(all_files)} in total") + sampled_files = np.random.choice( + all_files, replace=False, size=int(sample_ratio * len(all_files)) + ) + print(f"{len(sampled_files)} sampled files") + + os.chdir(work_dir) + with open(f"sampled_binaries.txt", "w") as f: + for fname in sampled_files: + fname = os.path.basename(fname) + f.write(fname + "\n") + + print("creating tar file") + os.chdir("binary_files/") + tar( + "-cf", + f"../sampled_binaries_{sample_ratio}.tar", + "-T", + "../sampled_binaries.txt", + ) + + +if __name__ == "__main__": + main() diff --git a/prediction-plugin/utils/util.py b/prediction-plugin/utils/util.py index 2f0866a..8b01106 100644 --- a/prediction-plugin/utils/util.py +++ b/prediction-plugin/utils/util.py @@ -1,8 +1,9 @@ import collections import gc import os -import psutil import sys + +import psutil import torch @@ -16,7 +17,7 @@ class cached_property(object): """ def __init__(self, func): - self.__doc__ = getattr(func, '__doc__') + self.__doc__ = getattr(func, "__doc__") self.func = func def __get__(self, obj, cls): @@ -29,10 +30,12 @@ def __get__(self, obj, cls): def memReport(): for obj in gc.get_objects(): if torch.is_tensor(obj): - print(type(obj), - obj.size(), - obj.element_size() * obj.nelement() / 1024 / 1024, - file=sys.stderr) + print( + type(obj), + obj.size(), + obj.element_size() * obj.nelement() / 1024 / 1024, + file=sys.stderr, + ) def cpuStats(): @@ -41,8 +44,8 @@ def cpuStats(): print(psutil.virtual_memory()) # physical memory usage pid = os.getpid() py = psutil.Process(pid) - memoryUse = py.memory_info()[0] / 2. ** 30 # memory use in GB...I think - print('memory GB:', memoryUse, file=sys.stderr) + memoryUse = py.memory_info()[0] / 2.0 ** 30 # memory use in GB...I think + print("memory GB:", memoryUse, file=sys.stderr) def update(d, u): diff --git a/prediction-plugin/utils/vocab.py b/prediction-plugin/utils/vocab.py index 61b27c4..4ce9d2c 100644 --- a/prediction-plugin/utils/vocab.py +++ b/prediction-plugin/utils/vocab.py @@ -10,18 +10,18 @@ --freq-cutoff= frequency cutoff [default: 2] """ +import json +import pickle from collections import Counter from itertools import chain -from docopt import docopt -import json import sentencepiece as spm - +import torch +from docopt import docopt from utils.grammar import Grammar - -SAME_VARIABLE_TOKEN = '' -END_OF_VARIABLE_TOKEN = '' +SAME_VARIABLE_TOKEN = "" +END_OF_VARIABLE_TOKEN = "" PAD_ID = 0 @@ -34,18 +34,17 @@ def __init__(self, subtoken_model_path=None): self.subtoken_model = spm.SentencePieceProcessor() self.subtoken_model.load(subtoken_model_path) - subtoken_model = \ - subtoken_model_path[:subtoken_model_path.rfind('.model')] - vocab_path = f'{subtoken_model}.vocab' + subtoken_model = subtoken_model_path[: subtoken_model_path.rfind(".model")] + vocab_path = f"{subtoken_model}.vocab" for i, line in enumerate(open(vocab_path)): - word, prob = line.strip().split('\t') + word, prob = line.strip().split("\t") self.word2id[word] = len(self.word2id) else: self.subtoken_model = None - self.word2id[''] = PAD_ID - self.word2id[''] = 1 - self.word2id[''] = 2 - self.word2id[''] = 3 + self.word2id[""] = PAD_ID + self.word2id[""] = 1 + self.word2id[""] = 2 + self.word2id[""] = 3 self.word2id[SAME_VARIABLE_TOKEN] = 4 self.id2word = {v: k for k, v in self.word2id.items()} @@ -55,7 +54,7 @@ def __getitem__(self, word): @property def unk_id(self): - return self.word2id[''] + return self.word2id[""] def is_unk(self, word): return word not in self.word2id @@ -64,13 +63,13 @@ def __contains__(self, word): return word in self.word2id def __setitem__(self, key, value): - raise ValueError('vocabulary is readonly') + raise ValueError("vocabulary is readonly") def __len__(self): return len(self.word2id) def __repr__(self): - return 'Vocabulary[size=%d]' % len(self) + return "Vocabulary[size=%d]" % len(self) def id2word(self, wid): return self.id2word[wid] @@ -85,34 +84,37 @@ def add(self, word): @property def params(self): - params = dict(unk_id=self.unk_id, word2id=self.word2id, - subtoken_model_path=self.subtoken_model_path) - if 'word_freq' in self.__dict__: - params['word_freq'] = self.word_freq + params = dict( + unk_id=self.unk_id, + word2id=self.word2id, + subtoken_model_path=self.subtoken_model_path, + ) + if "word_freq" in self.__dict__: + params["word_freq"] = self.word_freq return params def save(self, path): - json.dump(self.params, open(path, 'w'), indent=2) + json.dump(self.params, open(path, "w"), indent=2) @classmethod def load(cls, path=None, params=None): if path: - params = json.load(open(path, 'r')) + params = json.load(open(path, "r")) else: - assert params, 'Params must be given when path is None!' + assert params, "Params must be given when path is None!" - if 'subtoken_model_path' in params: - subtoken_model_path = params['subtoken_model_path'] + if "subtoken_model_path" in params: + subtoken_model_path = params["subtoken_model_path"] else: subtoken_model_path = None entry = VocabEntry(subtoken_model_path) - setattr(entry, 'word2id', params['word2id']) - setattr(entry, 'id2word', {v: k for k, v in params['word2id'].items()}) - if 'word_freq' in params: - setattr(entry, 'word_freq', params['word_freq']) + setattr(entry, "word2id", params["word2id"]) + setattr(entry, "id2word", {v: k for k, v in params["word2id"].items()}) + if "word_freq" in params: + setattr(entry, "word_freq", params["word_freq"]) return entry @@ -122,14 +124,14 @@ def from_corpus(corpus, size, freq_cutoff=0, predefined_tokens=None): word_freq = Counter(chain(*corpus)) freq_words = [w for w in word_freq if word_freq[w] >= freq_cutoff] - print(f'number of word types: {len(word_freq)}, ' - f'number of word types w/ frequency >= {freq_cutoff}: ' - f'{freq_words}') - top_k_words = sorted( - word_freq, - key=lambda x: (-word_freq[x], x) - )[:size] - print('top 30 words: %s' % ', '.join(top_k_words[:30])) + print( + f"number of word types: {len(word_freq)}, " + f"number of word types w/ frequency >= {freq_cutoff}: " + f"{freq_words}" + ) + + top_k_words = sorted(word_freq, key=lambda x: (-word_freq[x], x))[:size] + print("top 30 words: %s" % ", ".join(top_k_words[:30])) if predefined_tokens: for token in predefined_tokens: @@ -141,7 +143,7 @@ def from_corpus(corpus, size, freq_cutoff=0, predefined_tokens=None): vocab_entry.add(word) # store the work frequency table in the - setattr(vocab_entry, 'word_freq', word_freq) + setattr(vocab_entry, "word_freq", word_freq) return vocab_entry @@ -155,9 +157,10 @@ def __init__(self, **kwargs): self.entries.append(key) def __repr__(self): - words = ', '.join(f'{entry} {getattr(self, entry)}words' - for entry in self.entries) - return f'Vocab({words})' + words = ", ".join( + f"{entry} {getattr(self, entry)}words" for entry in self.entries + ) + return f"Vocab({words})" @property def params(self): @@ -168,14 +171,14 @@ def params(self): return params def save(self, path): - json.dump(self.params, open(path, 'w'), indent=2) + json.dump(self.params, open(path, "w"), indent=2) - @classmethod - def load(cls, path): - params = json.load(open(path, 'r')) + @staticmethod + def load(path): + params = json.load(open(path, "r")) entries = dict() for key, val in params.items(): - if key in ('grammar', ): + if key in ("grammar",): entry = Grammar.load(val) else: entry = VocabEntry.load(params=val) @@ -183,17 +186,17 @@ def load(cls, path): return cls(**entries) -if __name__ == '__main__': +if __name__ == "__main__": from utils.dataset import Dataset args = docopt(__doc__) - vocab_size = int(args['--size']) - vocab_file = args['VOCAB_FILE'] - train_set = Dataset(args['TRAIN_FILE']) + vocab_size = int(args["--size"]) + vocab_file = args["VOCAB_FILE"] + train_set = Dataset(args["TRAIN_FILE"]) - src_code_tokens_file = vocab_file + '.src_code_tokens.txt' + src_code_tokens_file = vocab_file + ".src_code_tokens.txt" src_preserved_tokens = set() - f_src_token = open(src_code_tokens_file, 'w') + f_src_token = open(src_code_tokens_file, "w") # extract vocab and node types node_types = set() @@ -214,95 +217,99 @@ def load(cls, path): if old_var_name != new_var_name: tgt_words.append(new_var_name) - if node.node_type == 'obj' \ - or node.node_type == 'block' \ - and hasattr(node, 'name'): + if ( + node.node_type == "obj" + or node.node_type == "block" + and hasattr(node, "name") + ): identifier_names.append(node.name) - if hasattr(node, 'type_tokens'): + if hasattr(node, "type_tokens"): type_tokens.extend(node.type_tokens) code_tokens = example.code_tokens - preserved_tokens = [token for token in code_tokens - if token.startswith('@@') and token.endswith('@@')] + preserved_tokens = [ + token + for token in code_tokens + if token.startswith("@@") and token.endswith("@@") + ] src_preserved_tokens.update(preserved_tokens) - f_src_token.write(' '.join(code_tokens) + '\n') + f_src_token.write(" ".join(code_tokens) + "\n") f_src_token.close() - print('building source words vocabulary') + print("building source words vocabulary") src_var_vocab_entry = VocabEntry.from_corpus( - [src_words], size=vocab_size, freq_cutoff=int(args['--freq-cutoff']) + [src_words], size=vocab_size, freq_cutoff=int(args["--freq-cutoff"]) ) - if args['--use-bpe']: - print('use bpe') + if args["--use-bpe"]: + print("use bpe") - print('building source code tokens vocabulary') + print("building source code tokens vocabulary") # train subtoken models - src_preserved_tokens = ','.join(src_preserved_tokens) + src_preserved_tokens = ",".join(src_preserved_tokens) spm.SentencePieceTrainer.Train( - '--add_dummy_prefix=false ' - f'--pad_id={PAD_ID} ' - '--bos_id=1 ' - '--eos_id=2 ' - '--unk_id=3 ' - f'--user_defined_symbols={src_preserved_tokens} ' - f'--vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.src_code_tokens ' - '--model_type=bpe ' - f'--input={src_code_tokens_file}' + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + f"--user_defined_symbols={src_preserved_tokens} " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.src_code_tokens " + "--model_type=bpe " + f"--input={src_code_tokens_file}" ) - src_code_tokens_vocab_entry = \ - VocabEntry(f'{vocab_file}.src_code_tokens.model') + src_code_tokens_vocab_entry = VocabEntry(f"{vocab_file}.src_code_tokens.model") - print('building target words vocabulary') - tgt_word_file = args['VOCAB_FILE'] + '.var_names.txt' - with open(tgt_word_file, 'w') as f: + print("building target words vocabulary") + tgt_word_file = args["VOCAB_FILE"] + ".var_names.txt" + with open(tgt_word_file, "w") as f: for name in tgt_words: - f.write(name + '\n') + f.write(name + "\n") spm.SentencePieceTrainer.Train( - '--add_dummy_prefix=false ' - f'--pad_id={PAD_ID} ' - '--bos_id=1 ' - '--eos_id=2 ' - '--unk_id=3 ' - '--control_symbols= ' - f'--vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.tgt ' - '--model_type=bpe ' - f'--input={tgt_word_file}' + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + "--control_symbols= " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.tgt " + "--model_type=bpe " + f"--input={tgt_word_file}" ) - tgt_var_vocab_entry = VocabEntry(f'{vocab_file}.tgt.model') + tgt_var_vocab_entry = VocabEntry(f"{vocab_file}.tgt.model") else: tgt_var_vocab_entry = VocabEntry.from_corpus( [tgt_words], size=vocab_size, - freq_cutoff=int(args['--freq-cutoff']), - predefined_tokens=[SAME_VARIABLE_TOKEN] + freq_cutoff=int(args["--freq-cutoff"]), + predefined_tokens=[SAME_VARIABLE_TOKEN], ) - id_names_file = vocab_file + '.id_names.txt' - with open(id_names_file, 'w') as f: + id_names_file = vocab_file + ".id_names.txt" + with open(id_names_file, "w") as f: for name in identifier_names: - f.write(name + '\n') + f.write(name + "\n") - print('train subtoken model for obj names') + print("train subtoken model for obj names") # train subtoken models spm.SentencePieceTrainer.Train( - '--add_dummy_prefix=false ' - f'--pad_id={PAD_ID} ' - '--bos_id=1 ' - '--eos_id=2 ' - '--unk_id=3 ' - '--control_symbols= ' - f'--vocab_size={vocab_size} ' - f'--model_prefix={vocab_file}.obj_name ' - '--model_type=bpe ' - f'--input={id_names_file}' + "--add_dummy_prefix=false " + f"--pad_id={PAD_ID} " + "--bos_id=1 " + "--eos_id=2 " + "--unk_id=3 " + "--control_symbols= " + f"--vocab_size={vocab_size} " + f"--model_prefix={vocab_file}.obj_name " + "--model_type=bpe " + f"--input={id_names_file}" ) - obj_name_vocab_entry = VocabEntry(f'{vocab_file}.obj_name.model') + obj_name_vocab_entry = VocabEntry(f"{vocab_file}.obj_name.model") type_vocab = Counter(type_tokens) num_types = 100 @@ -312,16 +319,18 @@ def load(cls, path): print(type_token, freq) var_types.append(type_token) - print('init node types and variable types') + print("init node types and variable types") grammar = Grammar(node_types, var_types) - print('Node types:', node_types) - print('Variable types:', var_types) + print("Node types:", node_types) + print("Variable types:", var_types) - vocab = Vocab(source=src_var_vocab_entry, - source_tokens=src_code_tokens_vocab_entry, - target=tgt_var_vocab_entry, - obj_name=obj_name_vocab_entry, - grammar=grammar) + vocab = Vocab( + source=src_var_vocab_entry, + source_tokens=src_code_tokens_vocab_entry, + target=tgt_var_vocab_entry, + obj_name=obj_name_vocab_entry, + grammar=grammar, + ) - vocab.save(args['VOCAB_FILE']) + vocab.save(args["VOCAB_FILE"])