<reponame>MirageLyu/code-prediction-transformer<filename>code2seq/code2seq_model.py
#!/usr/bin/env python3
# Copyright (c) 2019 Technion
# Copyright (c) Facebook, Inc. and its affiliates.
# ----------------------------------------------------------------------------
# MIT License

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ----------------------------------------------------------------------------
"""
    Code2seq model is adapted from: https://github.com/tech-srl/code2seq
"""

import torch
import torch.nn as nn
import torch.nn.functional as F



class EmbeddingAttentionLayer(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.dim = dim
        self.attention = torch.randn(1, dim)
        self.attention = nn.Parameter(self.attention)

    def compute_weights(self, embedded: torch.Tensor) -> torch.Tensor:
        unnormalized_weights = embedded.matmul(self.attention.t())
        attention_weights = F.softmax(unnormalized_weights, dim=1)
        return attention_weights

    def forward(self, embedded: torch.Tensor) -> torch.Tensor:
        attention_weights = self.compute_weights(embedded)
        weighted = torch.bmm(attention_weights.transpose(1, 2), embedded)
        return weighted


class Code2SeqModel(nn.Module):
    def __init__(
        self,
        token_vocab_size: int,
        subtoken_vocab_size: int,
        output_vocab_size: int,
        token_pad_idx: int,
        subtoken_pad_idx: int,
        loss_fn: nn.Module,
        n_embd: int = 128,
        rnn_dropout: float = 0.5,
        embed_dropout: float = 0.25,
    ):
        super().__init__()
        self.subtoken_embedding = nn.Embedding(
            subtoken_vocab_size, n_embd, padding_idx=subtoken_pad_idx
        )
        self.node_embedding = nn.Embedding(
            token_vocab_size, n_embd, padding_idx=token_pad_idx
        )
        self.path_lstm = nn.LSTM(
            n_embd, n_embd, bidirectional=True, dropout=rnn_dropout, batch_first=True
        )
        self.combined_layer = nn.Linear(n_embd * 4, n_embd)
        self.dropout = nn.Dropout(embed_dropout)
        self.attn_layer = EmbeddingAttentionLayer(n_embd)
        self.out_layer = nn.Linear(n_embd, output_vocab_size)

        self.loss_fn = loss_fn

    def reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def embed_paths(self, paths):
        path_tokens_embedded = self.node_embedding(paths)
        batch_size, bag_size, path_len, _ = path_tokens_embedded.shape
        path_tokens_embedded = path_tokens_embedded.view(
            (batch_size * bag_size, path_len, -1)
        )
        out, (h, c) = self.path_lstm(path_tokens_embedded)
        paths_embedded = h.permute((1, 0, 2)).reshape(batch_size, bag_size, -1)
        return paths_embedded

    def embed_subtokens(self, subtokens):
        tokens_embedded = self.subtoken_embedding(subtokens)
        return tokens_embedded.sum(2)

    def forward(self, starts, paths, ends, targets, return_loss=False):
        # embed individual parts
        starts_embedded = self.embed_subtokens(starts)
        paths_embedded = self.embed_paths(paths)
        ends_embedded = self.embed_subtokens(ends)

        # combine by concacenating
        combined_embedded = torch.cat(
            (starts_embedded, paths_embedded, ends_embedded), dim=2
        )
        combined_embedded = self.dropout(combined_embedded)
        combined_embedded = torch.tanh(self.combined_layer(combined_embedded))

        # combine paths by simple attention
        code_embedded = self.attn_layer(combined_embedded).squeeze()
        y_pred = self.out_layer(code_embedded)

        if not return_loss:
            return y_pred
        return self.loss_fn(y_pred, targets)

#!/usr/bin/env python3
# Copyright (c) 2019 Technion
# Copyright (c) Facebook, Inc. and its affiliates.
# ----------------------------------------------------------------------------
# MIT License

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ----------------------------------------------------------------------------
"""
    Code2seq model is adapted from: https://github.com/tech-srl/code2seq
"""

import torch
import torch.nn as nn
import torch.nn.functional as F



class EmbeddingAttentionLayer(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.dim = dim
        self.attention = torch.randn(1, dim)
        self.attention = nn.Parameter(self.attention)

    def compute_weights(self, embedded: torch.Tensor) -> torch.Tensor:
        unnormalized_weights = embedded.matmul(self.attention.t())
        attention_weights = F.softmax(unnormalized_weights, dim=1)
        return attention_weights

    def forward(self, embedded: torch.Tensor) -> torch.Tensor:
        attention_weights = self.compute_weights(embedded)
        weighted = torch.bmm(attention_weights.transpose(1, 2), embedded)
        return weighted


class Code2SeqModel(nn.Module):
    def __init__(
        self,
        token_vocab_size: int,
        subtoken_vocab_size: int,
        output_vocab_size: int,
        token_pad_idx: int,
        subtoken_pad_idx: int,
        loss_fn: nn.Module,
        n_embd: int = 128,
        rnn_dropout: float = 0.5,
        embed_dropout: float = 0.25,
    ):
        super().__init__()
        self.subtoken_embedding = nn.Embedding(
            subtoken_vocab_size, n_embd, padding_idx=subtoken_pad_idx
        )
        self.node_embedding = nn.Embedding(
            token_vocab_size, n_embd, padding_idx=token_pad_idx
        )
        self.path_lstm = nn.LSTM(
            n_embd, n_embd, bidirectional=True, dropout=rnn_dropout, batch_first=True
        )
        self.combined_layer = nn.Linear(n_embd * 4, n_embd)
        self.dropout = nn.Dropout(embed_dropout)
        self.attn_layer = EmbeddingAttentionLayer(n_embd)
        self.out_layer = nn.Linear(n_embd, output_vocab_size)

        self.loss_fn = loss_fn

    def reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def embed_paths(self, paths):
        path_tokens_embedded = self.node_embedding(paths)
        batch_size, bag_size, path_len, _ = path_tokens_embedded.shape
        path_tokens_embedded = path_tokens_embedded.view(
            (batch_size * bag_size, path_len, -1)
        )
        out, (h, c) = self.path_lstm(path_tokens_embedded)
        paths_embedded = h.permute((1, 0, 2)).reshape(batch_size, bag_size, -1)
        return paths_embedded

    def embed_subtokens(self, subtokens):
        tokens_embedded = self.subtoken_embedding(subtokens)
        return tokens_embedded.sum(2)

    def forward(self, starts, paths, ends, targets, return_loss=False):
        # embed individual parts
        starts_embedded = self.embed_subtokens(starts)
        paths_embedded = self.embed_paths(paths)
        ends_embedded = self.embed_subtokens(ends)

        # combine by concacenating
        combined_embedded = torch.cat(
            (starts_embedded, paths_embedded, ends_embedded), dim=2
        )
        combined_embedded = self.dropout(combined_embedded)
        combined_embedded = torch.tanh(self.combined_layer(combined_embedded))

        # combine paths by simple attention
        code_embedded = self.attn_layer(combined_embedded).squeeze()
        y_pred = self.out_layer(code_embedded)

        if not return_loss:
            return y_pred
        return self.loss_fn(y_pred, targets)
#!/usr/bin/env python3 # Copyright (c) 2019 Technion # Copyright (c) Facebook, Inc. and its affiliates. # ---------------------------------------------------------------------------- # MIT License # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # ---------------------------------------------------------------------------- Code2seq model is adapted from: https://github.com/tech-srl/code2seq # embed individual parts # combine by concacenating # combine paths by simple attention