Skip to content

Refactor duplicate code #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
probinso wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Refactor duplicate code #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
218 changes: 135 additions & 83 deletions neural-model/model/attentional_recurrent_subtoken_decoder.py
View file
Open in desktop

Large diffs are not rendered by default.

41 changes: 28 additions & 13 deletions neural-model/model/embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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):
Expand Down Expand Up @@ -55,29 +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 = embedding.sum(dim=1) / (
sub_tokens_mask.sum(-1) + SMALL_NUMBER
).unsqueeze(-1)

return embedding
1 change: 0 additions & 1 deletion neural-model/model/encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,3 @@

class Encoder(nn.Module):
pass

177 changes: 120 additions & 57 deletions neural-model/model/ensemble_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
}

Expand All @@ -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 = []
Expand All @@ -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)
Expand All @@ -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 = []
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)

Expand Down
Loading