seq2seq2.py

# -*- coding: utf-8 -*-
"""Seq2Seq.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Xnp8K05GDzNwc3XZMQTWq3ZVzN_tUmjy

# Seq2Seq with Attentionを実装する

## 大域変数の定義

データセットやパラメータを定義する  
大域変数は全て大文字で表記することとする
"""

# DEFINE GLOBAL VARIABLE
TRAIN_FILE='data/train.tsv'
VALID_FILE='data/valid.tsv'
TEST_FILE='data/test.tsv'
MAKE_DATA=True
BATCH_SIZE=10
DEVICE='cuda'
EPOCH=10
EMBED=256
HIDDEN=256
PARALLEL=True
DEVICES=[0,1]

import os 
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

"""## 対訳データセットのアップロード

今回の実装では[Tanaka Corpus](http://www.edrdg.org/wiki/index.php/Tanaka_Corpus)を使用する  
田中コーパスは大学生らが集めた日英の対訳データセットとなっている  
以下のプログラムを実行し、`ファイルを選択`ボタンからローカルに用意したコーパスをアップロードする

ここでデータをアップロードせずに、
**データセットの保存**から開始すると10件のみのデータセットで動作を確認できる
"""

# Upload Dataset
#from google.colab import files
#uploaded = files.upload()

"""## データのクリーニング

pythonスクリプトを用いて、データの前処理を行う  
田中コーパスは以下の形式で記述されている

    A: ムーリエルは２０歳になりました。	Muiriel is 20 now.#ID=1282_4707
    B: は 二十歳(はたち){２０歳} になる[01]{になりました}
    A: すぐに戻ります。	I will be back soon.#ID=1284_4709
    B: 直ぐに{すぐに} 戻る{戻ります}

したがって、まず、`A: `から始まる文を抽出する  
次に、`\t`で分割を行い、`#ID`以降の余分な情報を取り除く
"""

# Read Dataset
import gzip
#import copy
if MAKE_DATA:
    with gzip.open('data/examples.utf.gz', 'rt') as f:
        dataset = [line.strip().split('\t') for line in f.readlines() if line.startswith('A')]
    print('Dataset size: ', len(dataset))
    print('raw dataset')
    for example in dataset[:5]:
        print(example)

    # Cleaning Text
    for i in range(len(dataset)):
        src = dataset[i][0]
        trg = dataset[i][1]
        src = src.split('A: ')[1]
        trg = trg.split('#ID=')[0]
        dataset[i][0] = src
        dataset[i][1] = trg

    print('clening dataset')
    for example in dataset[:5]:
        print(example)

    """## データセットの保存
    データセットを`train`、`validation`、`test`に分割して保存する
    """    
    # Write out Dataset
    import random
    if len(dataset) == 10: # mini dataset
        train = dataset[:6]
        valid = dataset[6:8]
        test  = dataset[8:]
    else:
        random.seed(0)
        random.shuffle(dataset)
        train = dataset[     0:100000]
        valid = dataset[100000:105000]
        test  = dataset[110000:115000]

    def write_tsv(file_name, dataset):
        text = '\n'.join(['\t'.join(pair) for pair in dataset])
        with open(file_name, 'w') as f:
            print(text, file=f)
        
    write_tsv(TRAIN_FILE, train)
    write_tsv(VALID_FILE, valid)
    write_tsv(TEST_FILE, test)

"""## Mecab, spaCyのインストール
日本語の分かち書きに`Mecab`、英語のtokenizerに`spaCy`を使用するためにインストールを行う
"""

# Install Mecab and Spacy (tokenizer)
"""
!echo 'install mecab'
!apt-get -q -y install sudo file mecab libmecab-dev mecab-ipadic-utf8 git curl python-mecab > /dev/null
!pip install mecab-python3 > /dev/null
!echo 'install spaCy'
!pip install spaCy > /dev/null
!python -m spacy download en > /dev/null
!echo 'complete.'
"""
"""spaCyの使用例を以下に示す"""

# Spacy Sample
import spacy
nlp = spacy.load('en')
doc = nlp('Hello, world. Here are two sentences.')
example = [str(token) for token in doc]
print(example)

"""Mecabの使用例を以下に示す"""

# Mecab Sample
import MeCab
mecab = MeCab.Tagger('-Owakati')
example = mecab.parse('こんにちは、世界。これは二文目です。')
print(example)

"""## torch,torchtextのインストール"""
"""
!pip install torch
!pip install git+https://github.com/pytorch/text/
"""
import torch
#import torchtext

"""## torchtextを用いたデータの読み込み"""

from torchtext import data

# Prepare Data Fields
src_TEXT = data.Field(
    sequential=True, include_lengths=True,
    init_token='<sos>', eos_token='<eos>',
    lower=True, tokenize=lambda s: mecab.parse(s).rstrip().split())
trg_TEXT = data.Field(
    sequential=True, include_lengths=True,
    init_token='<sos>', eos_token='<eos>',
    lower=True, tokenize='spacy')
fields = [('src', src_TEXT), ('trg', trg_TEXT)]

# Load Dataset Using torchtext
train, val, test = data.TabularDataset.splits(
    path='./',
    train=TRAIN_FILE,
    validation=VALID_FILE,
    test=TEST_FILE,
    format='tsv',
    fields=fields)

# Build Vocablary
src_TEXT.build_vocab(train, min_freq=2)
trg_TEXT.build_vocab(train, min_freq=2)
source_vocab_size = len(src_TEXT.vocab.itos)
target_vocab_size = len(trg_TEXT.vocab.itos)

# Make iterator
train_iter, valid_iter = data.Iterator.splits(
    (train, val), batch_size=BATCH_SIZE, repeat=False,
    sort_key=lambda x: data.interleave_keys(len(x.src), len(x.trg)),
    device=torch.device(DEVICE))

test_iter = data.Iterator(
    test, batch_size=BATCH_SIZE, repeat=False, device=torch.device(DEVICE))

print(torch.device(DEVICE))

# Check Dataset
print(train[0].src)
print(train[0].trg)
# Check Vocab
print('src vocab: ', len(src_TEXT.vocab.itos))
print(src_TEXT.vocab.itos[:10])
print('trg vocab: ', len(trg_TEXT.vocab.itos))
print(trg_TEXT.vocab.itos[:10])

"""## Seq2Seq with Att モデル

* Seq2Seqクラス
* Encoderクラス
* Decoderクラス
* Generatorクラス
* Attentionクラス

から構成される。

### 各クラスの関係
Seq2SeqがEncoder, Deocder, Generatorを持つ。  
DecoderはAttentionを持つ。

### 各クラスの役割
Encoderは入力文を圧縮してベクトルへ変換する。  
Decoderは逐次の語彙次元へ写像する直前の隠れ状態を生成。  
GeneratorがDecoderの生成した隠れ状態を語彙次元へと写像。
Attentionはデコードの各時刻にエンコーダの状態を参照するために用いられる。
"""

import torch.nn as nn
import torch.nn.functional as F
import numpy
# Define Seq2Seq with Attention Model
class Seq2Seq(nn.Module):
    def __init__(self,
                 src_vocab_size, trg_vocab_size,
                 embed_size, hidden_size, dropout_p):
        super(Seq2Seq, self).__init__()
        self.encoder = Encoder(src_vocab_size, embed_size, hidden_size, dropout_p)
        self.decoder = Decoder(trg_vocab_size, embed_size, hidden_size, dropout_p)
        self.generator = Generator(trg_vocab_size, hidden_size, dropout_p)

    def forward(self, src, trg):
        if len(src[1].size()) == 2:
            # for compatibility of parallel
            src = (src[0], src[1].squeeze(0))

        return self.decoder(trg, *self.encoder(src), src[0] == 1)

class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, dropout_p):
        super(Encoder, self).__init__()
        self.hidden_size = hidden_size
        self.n_layers = 2
        self.embed = nn.Embedding(vocab_size, embed_size, padding_idx=1)
        self.gru = nn.GRU(embed_size, hidden_size, num_layers=self.n_layers, 
                                     dropout=dropout_p, bidirectional=True)
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, src):
        # 単語と文長に分解
        words = src[0] #LxB
        lengths = src[1] #B
        total_length=words.size()[0]

        # 入力を文長で降順ソート
        lengths, parm_idx = torch.sort(lengths, 0, descending=True)
        device = parm_idx.device
        parm_idx_rev = torch.tensor(_inverse_indices(parm_idx), device=device)
        words = words[:, parm_idx]

        # Embedding
        ex = self.embed(words) #LxB -> LxBxH
        ex = self.dropout(ex)

        # 入力をGRUで処理するためにPackする
        packed_input = nn.utils.rnn.pack_padded_sequence(ex, lengths)

        # GRU (L 回処理される)
        packed_output, hidden = self.gru(packed_input)

        # GRUの出力をUnPackしてPaddingする
        output = nn.utils.rnn.pad_packed_sequence(packed_output, total_length=total_length)

        # ソートを元の順へ戻す
        output = output[0][:, parm_idx_rev]  #LxBx2H
        hidden = hidden[:,parm_idx_rev]
	
        # 前向きと後ろ向きの平均を計算
        output = output[:, :, :self.hidden_size] + output[:, :, self.hidden_size:]
	
        return output, hidden

class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, dropout_p):
        super(Decoder, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size, padding_idx=1)
        self.n_layers = 2
        self.gru = nn.GRU(embed_size, hidden_size,
                          num_layers=self.n_layers, dropout=dropout_p)
        self.Wc = nn.Linear(2*hidden_size, hidden_size)
        self.attention = Attention()
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, trg, encoder_states, last_states, src_mask):
        words   = trg[0][:-1] #LxB (Lは<eos>を抜いた長さ)
        last_states = last_states[:self.n_layers]

        #encoder_states # LxBxH
        #last_states[0], [1] #n_layersxBxH

        output = []
        for word in words: #一単語ずつ処理する
            out, context, last_state = self.forward_step(
                word, encoder_states, last_states, src_mask)
            output.append(out)

        return output

    def forward_step(self, word, encoder_states, states, src_mask):
        # Embedding
        ex = self.embed(word) #BxH
        ex = self.dropout(ex)

        # input-feed
        rnn_input = ex #Bx2H
        # GRUの入力に文長の次元がいるため拡張
        rnn_input = torch.unsqueeze(rnn_input, 0) #1xBx2H
        # LSTM(1単語ずつ)
        rnn_output, status = self.lstm(rnn_input, states)
        # 文脈ベクトルを計算
        context = self.attention(rnn_output, encoder_states, src_mask)
        # 拡張した次元を元に戻す
        rnn_output = torch.squeeze(rnn_output, 0)
        t = torch.cat((rnn_output, context), 1) #Bx2H
        # 次元を元に戻す(2H->H)
        t = self.dropout(t)
        out = self.Wc(t) #BxH
        out = torch.tanh(self.Wc(t)) #Bx1xH
        return out, context, status

class Generator(nn.Module):
    def __init__(self, vocab_size, hidden_size, dropout_p):
        super(Generator, self).__init__()
        self.Wo = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, decoder_states):
        if type(decoder_states) is list:
            decoder_states = torch.stack(decoder_states, 0) #LxBxH
        # 語彙サイズのベクトルへと写像
        decoder_states = self.dropout(decoder_states)
        out = self.Wo(decoder_states)
        return out

class Attention(nn.Module):
    def __init__(self):
        super(Attention, self).__init__()

    def forward(self, decoder_state, encoder_states, src_mask):
        #decoder_state  1xBxH
        #encoder_states LxBxH
        decoder_state = decoder_state.transpose(0, 1) #Bx1xH
        decoder_state = decoder_state.transpose(1, 2) #BxHx1
        encoder_states = encoder_states.transpose(0, 1) #BxLxH
        # 内積でスコアを計算
        score = torch.matmul(encoder_states, decoder_state) #BxLx1
        # Paddingに当たる部分の値を-700へ(softmaxで0に落ちるように)
        src_mask = src_mask.transpose(0, 1).unsqueeze(2) #BxLx1
        score = torch.where(src_mask, torch.full_like(score, -700), score)
        # softmaxで確率化
        weight = F.softmax(score, dim=1) # BxLx1
        context = torch.matmul(weight.transpose(1,2), encoder_states).squeeze(1)
        return context

def _inverse_indices(indices):
    indices = indices.cpu().numpy()
    r = numpy.empty_like(indices)
    r[indices] = numpy.arange(len(indices))
    return r

"""## インスタンスの定義

モデル、オプティマイザなどの定義
"""

# Set up Model
model = Seq2Seq(source_vocab_size, target_vocab_size, EMBED, HIDDEN, 0.2)

print(torch.device(DEVICE))
model.to(device=torch.device(DEVICE))

# Set up Optimizer
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=1)

# Define loss_compute
class LossCompute():
    def __init__(self, generator, criterion, opt=None):
        self.generator = generator
        self.criterion = criterion
        self.opt = opt

    def __call__(self, x, t, norm=None, backward=True):
        y = self.generator(x)

        loss = self.criterion(
            y.contiguous().view(-1, y.size(-1)), 
            t[1:].contiguous().view(-1))

        if norm is not None:
            loss = loss

        if backward:
            loss.backward()

        if self.opt is not None:
            self.opt.step()
            self.opt.zero_grad()

        return loss.item()

# Define parallel_loss_compute
class ParallelLossCompute():
    def __init__(self, generator, criterion, device_ids, opt=None):
        self.generator = generator
        self.criterion = nn.parallel.replicate(criterion, device_ids)
        self.device_ids = device_ids
        self.opt = opt

    def __call__(self, outputs, targets, norm=None, backward=True, dim=1):
        outputs = torch.stack(outputs, 0)
        output_device = self.device_ids[0]
        generator = nn.parallel.replicate(self.generator, self.device_ids)
        outputs = nn.parallel.scatter(outputs, self.device_ids, dim)
        outputs = [(output, ) for output in outputs]
        gen = nn.parallel.parallel_apply(generator, outputs)
        targets = nn.parallel.scatter(targets, self.device_ids, dim)
        # Compute loss
        y = [(g.contiguous().view(-1, g.size(-1)), t[1:].contiguous().view(-1))
             for g,t in zip(gen, targets)]
        loss = nn.parallel.parallel_apply(self.criterion, y)
        loss = [l.unsqueeze(0) for l in loss]
        loss = nn.parallel.gather(loss, output_device).mean()
        if norm is not None:
            loss = loss

        if backward:
            loss.backward()

        if self.opt is not None:
            self.opt.step()
            self.opt.zero_grad()

        return loss.item()

def data_parallel(module, inputs, device_ids, dim=1):
    if not device_ids:
        return module(inputs)
    
    output_device = device_ids[0]
    replicas = nn.parallel.replicate(module, device_ids)
    # dim length expand for scatter
    inputs = [(inputs[0][0], inputs[0][1].unsqueeze(0)),
              (inputs[1][0], inputs[1][1].unsqueeze(0))]
    inputs = nn.parallel.scatter(inputs, device_ids, dim)
    replicas = replicas[:len(inputs)]
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    return nn.parallel.gather(outputs, output_device)
    
"""## trainとvalidの関数化

1エポック分の処理を関数化する
"""

# Training Step
def train_step(iterator, model, loss_compute):
    model.train()
    temp_loss, total_loss = 0, 0
    for batch in iterator:
        if iterator.iterations % 100 == 0:
            print('iteration {}\t loss {}'.format(iterator.iterations, temp_loss / 100))
            temp_loss = 0
        if PARALLEL:
            out = data_parallel(model, (batch.src, batch.trg), device_ids=DEVICES, dim=1)
            loss = loss_compute(out, batch.trg[0], norm=batch.batch_size)
        else:
            out = model(batch.src, batch.trg)
            loss = loss_compute(out, batch.trg[0], norm=batch.batch_size)
        total_loss += loss
        temp_loss += loss
    return total_loss / iterator.iterations

# Validate Step
def valid_step(iterator, model, loss_compute):
    model.eval()
    total_loss = 0
    for batch in iterator:
        if PARALLEL:
            out = data_parallel(model, (batch.src, batch.trg), device_ids=DEVICES, dim=1)
            loss = loss_compute(out, batch.trg[0], norm=batch.batch_size, backward=False)
        else:
            out = model(batch.src, batch.trg)
            with torch.no_grad():
                loss = loss_compute(out, batch.trg[0],
                                    norm=batch.batch_size, backward=False)
        total_loss += loss
    return total_loss / iterator.iterations

"""## モデルの学習

モデルを学習する  
毎エポックでtrain_stepとvalid_stepを繰り返す。
"""

# Train loop
for epoch in range(EPOCH):
    print('epoch', epoch)
    if PARALLEL:
        train_loss = train_step(
            train_iter, model, ParallelLossCompute(
                model.generator, criterion, DEVICES, optimizer))
        valid_loss = valid_step(
            valid_iter, model, ParallelLossCompute(
                model.generator, criterion, DEVICES, opt=None))
    else:
        train_loss = train_step(train_iter, model,
                                LossCompute(model.generator, criterion, optimizer))
        valid_loss = valid_step(valid_iter, model,
                                LossCompute(model.generator, criterion, None))
    print('epoch loss train {}\t valid {}'.format(train_loss, valid_loss))

"""## TEST

テストデータに対して文生成を行う。
Beem search and Greedy
"""
def beam_search(model, src, max_len, trg_vocab):
     sentences = []
     a = 0.6 #a is the 
     topvalue = 1
     model.eval()

     # <sos>トークンをbatch_size分用意する
     sos_label = trg_vocab.stoi['<sos>']
     eos_label = trg_vocab.stoi['<eos>']
     sos = torch.ones(1, 1).fill_(sos_label).type_as(src[0]) #1xB
     eos = torch.ones(1, 1).fill_(eos_label).type_as(src[0]) #1xB

     for i in range(src[0].size()[1]):
         k = 3 #set k for beam search
         # Encoder
         source = src[0].transpose(0,1)[i].view(-1,1)
         length = src[1][i].view(1)
         encoder_states, last_states = model.encoder((source,length))

         sequences = []
         result = []
         for j in range(k):
             sequences.append([sos, torch.ones(1).type_as(src[0]).float()])

         for di in range(max_len-1):
             if k != 0:
                 all_candidates = []
                 for row in range(k):
                     ys = sequences[row][0]
                     score = sequences[row][1]
                     out, attn_weights = model.decoder(
                           (torch.cat([ys, eos], 0), None), encoder_states, last_states, source == 1)
                     prob = model.generator(out).squeeze(2) #LxkxVocab
                     prob = F.softmax(prob,dim=2).transpose(0,1)
                     for count in range(k):
                         word = prob.topk(k,dim=2)[1][0][-1][count].view(-1,1)
                         word = torch.cat([ys, word], dim=0)

                         probability = torch.neg(torch.log(prob.topk(k,dim=2)[0][0][-1][count]))
                         #length normalization
                         lp_y = ((5 + len(word)) ** a) / ((5 + 1) ** a)
                         score = (torch.mul(score, probability) ) / lp_y
                         all_candidates.append([word, score])
                     if di == 0:
                         break

                 # order all candidates by score 从小到大排序
                 ordered = sorted(all_candidates, key=lambda tup:tup[topvalue].item())

                 # select k best
                 sequences = ordered[:k]
                 input_indexes = []
                 minus = 0
                 reset = []
                 for row in range(k):
                     if sequences[row][0][-1] == eos_label:
                         minus += 1
                         result.append(sequences[row])
                     else:
                         reset.append(sequences[row])
                
                 k = k - minus
                 sequences = reset

             else:
                 break

         if sequences != []:
             result = result + sequences

         ordered = sorted(result, key=lambda tup :tup[topvalue].item())
         ys = ordered[0][0][1:].squeeze(1).cpu().numpy() #LxB -> BxL
         sentence = label2word(ys, trg_vocab.itos)
         sentences.append(sentence)
     return sentences
	 
	 
# Greedy Decode
def greedy_decode(model, src, max_len, trg_vocab):
    model.eval()
    # Encoder
    encoder_states, last_states = model.encoder(src)

    # <sos>トークンをbatch_size分用意する
    sos_label = trg_vocab.stoi['<sos>']
    eos_label = trg_vocab.stoi['<eos>']
    batch_size = src[0].size(1)
    sos = torch.ones(1, batch_size).fill_(sos_label).type_as(src[0]) #1xB
    eos = torch.ones(1, batch_size).fill_(eos_label).type_as(src[0]) #1xB
    ys = sos

    for i in range(max_len-1):
        out = model.decoder(
               (torch.cat([ys, eos], 0), None), encoder_states, last_states, src[0] == 1)
        prob = model.generator(out)
        _, next_word = torch.max(prob, dim = 2)
        ys = torch.cat([sos, next_word], dim=0)
    ys = ys[1:].transpose(0,1).cpu().numpy() #LxB -> BxL
    sentences = []
    for line in ys:
        sentence = label2word(line, trg_vocab.itos)
        sentences.append(sentence)
    return sentences

def label2word(labels, vocab_itos):
    sentence = [vocab_itos[label] for label in labels]
    sentence = ' '.join(sentence)
    sentence = sentence.split('<pad>')[0]
    sentence = sentence.split('<eos>')[0]
    return sentence

test_sentences = []
for batch in test_iter:
    src = batch.src
    source = [label2word(line, src_TEXT.vocab.itos)
              for line in src[0].transpose(0,1).cpu().numpy()]
    with torch.no_grad():
        translation = greedy_decode(model, src, max_len=30,
                                    trg_vocab=trg_TEXT.vocab)
    test_sentences.extend(zip(source, translation))
    break

for sentence_pair in test_sentences[:10]:
    print('\t'.join(sentence_pair))

"""## モデルのセーブ、ロード
モデルをセーブしておくと、続きからの学習やテストに利用できる。
optimizerがAdamの場合は、optimizerも同時にセーブしておく。
"""

def save_model(save_path, model, optimizer):
    status = {
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict()
    }
    torch.save(status, save_path)

def load_model(load_path, model, optimizer):
    check_point = torch.load(load_path)
    model.load_state_dict(check_point['model'])
    optimizer.load_state_dict(check_point['optimizer'])

"""手元にモデルをダウンロードすることも可能"""

#セーブ
save_model('train_model.cp', model, optimizer)

#ダウンロード
#from google.colab import files
#files.download('train_model.cp')