Gru

In [ ]:

import torch
import torch.nn as nn
pad_token = 0
start_token = 1
end_token = 2  
class Encoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, num_hiddens, num_layers, direction=1):
        super(Encoder, self).__init__()
        self.direction = direction
        self.bidirectional = direction == 2
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.layers = nn.GRU(embedding_size, num_hiddens, num_layers, batch_first=True, bidirectional=self.bidirectional)

    def forward(self, x):
        x = self.embedding(x)
        return self.layers(x)

import torch import torch.nn as nn pad_token = 0 start_token = 1 end_token = 2 class Encoder(nn.Module): def __init__(self, vocab_size, embedding_size, num_hiddens, num_layers, direction=1): super(Encoder, self).__init__() self.direction = direction self.bidirectional = direction == 2 self.embedding = nn.Embedding(vocab_size, embedding_size) self.layers = nn.GRU(embedding_size, num_hiddens, num_layers, batch_first=True, bidirectional=self.bidirectional) def forward(self, x): x = self.embedding(x) return self.layers(x)

In [ ]:

class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, num_hiddens, num_layers, direction=1):
        super(Decoder, self).__init__()
        self.direction = direction
        self.bidirectional = direction == 2
        self.vocab_size = vocab_size
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.layers = nn.GRU(embedding_size, num_hiddens, num_layers, batch_first=True, bidirectional=self.bidirectional)
        self.fc = nn.Linear(direction * num_hiddens, vocab_size)

    def init_state(self, enc_outputs):
        state = enc_outputs[1]
        return state

    def forward(self, x, state):
        x = self.embedding(x)
        x, state = self.layers(x, state)
        x = self.fc(x)
        return x, state

class Decoder(nn.Module): def __init__(self, vocab_size, embedding_size, num_hiddens, num_layers, direction=1): super(Decoder, self).__init__() self.direction = direction self.bidirectional = direction == 2 self.vocab_size = vocab_size self.num_hiddens = num_hiddens self.num_layers = num_layers self.embedding = nn.Embedding(vocab_size, embedding_size) self.layers = nn.GRU(embedding_size, num_hiddens, num_layers, batch_first=True, bidirectional=self.bidirectional) self.fc = nn.Linear(direction * num_hiddens, vocab_size) def init_state(self, enc_outputs): state = enc_outputs[1] return state def forward(self, x, state): x = self.embedding(x) x, state = self.layers(x, state) x = self.fc(x) return x, state

In [ ]:

class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x, y):
        encoded_state = self.encoder(x)
        decoder_state = self.decoder.init_state(encoded_state)
        decoded = self.decoder(y, decoder_state)
        return decoded 

class EncoderDecoder(nn.Module): def __init__(self, encoder, decoder): super(EncoderDecoder, self).__init__() self.encoder = encoder self.decoder = decoder def forward(self, x, y): encoded_state = self.encoder(x) decoder_state = self.decoder.init_state(encoded_state) decoded = self.decoder(y, decoder_state) return decoded

In [ ]:

class GRU(EncoderDecoder):
    def __init__(self, vocab_size, num_hiddens, embedding_size, encoder_layers, decoder_layers, direction=1):
        encoder = Encoder(vocab_size, embedding_size, num_hiddens, encoder_layers, direction=direction)
        decoder = Decoder(vocab_size, embedding_size, num_hiddens, decoder_layers, direction=direction)
        super().__init__(encoder, decoder)
    
    def forward(self, x, y):
        return super().forward(x, y)

    def predict(self, source, source_len, max_len, device):
        pad_token = 0
        bos_token = 1
        eos_token = 2
        self.eval()

        source = source.to(device)
        encoded = self.encoder(source)
        state = self.decoder.init_state(encoded)
        
        tmp = torch.tensor([bos_token], dtype=torch.long).unsqueeze(0).to(device)
        preds = []
        for _ in range(max_len):
            out, _ = self.decoder(tmp, state)
            # 找到概率最大的
            out = out.argmax(dim=2)
            # 拼接
            tmp = torch.cat((tmp, out[:, -1].unsqueeze(0)), dim=1)
            # 获取模型预测的最后一步的标记
            end = tmp[:, -1].unsqueeze(0).item()
            if end == eos_token:
                break
            preds.append(end)
        return preds

class GRU(EncoderDecoder): def __init__(self, vocab_size, num_hiddens, embedding_size, encoder_layers, decoder_layers, direction=1): encoder = Encoder(vocab_size, embedding_size, num_hiddens, encoder_layers, direction=direction) decoder = Decoder(vocab_size, embedding_size, num_hiddens, decoder_layers, direction=direction) super().__init__(encoder, decoder) def forward(self, x, y): return super().forward(x, y) def predict(self, source, source_len, max_len, device): pad_token = 0 bos_token = 1 eos_token = 2 self.eval() source = source.to(device) encoded = self.encoder(source) state = self.decoder.init_state(encoded) tmp = torch.tensor([bos_token], dtype=torch.long).unsqueeze(0).to(device) preds = [] for _ in range(max_len): out, _ = self.decoder(tmp, state) # 找到概率最大的 out = out.argmax(dim=2) # 拼接 tmp = torch.cat((tmp, out[:, -1].unsqueeze(0)), dim=1) # 获取模型预测的最后一步的标记 end = tmp[:, -1].unsqueeze(0).item() if end == eos_token: break preds.append(end) return preds