DeepMig.py

# -*- coding: utf-8 -*-
"""DeepMig

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/14Vy6roEeHxKjXrjnRzfT061MyzADhvtT
"""

import pathlib
import random
import string
import re
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from datetime import datetime
#from tensorflow.keras.layers import TextVectorization
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization

from os.path import join

def decode_sequence(input_sentence):
    tokenized_input_sentence = input_vectorization([input_sentence])
    decoded_sentence = "[bos]"
    for i in range(max_decoded_sentence_length):
        tokenized_target_sentence = output_vectorization([decoded_sentence])[:, :-1]
        predictions = transformer([tokenized_input_sentence, tokenized_target_sentence])

        sampled_token_index = np.argmax(predictions[0, i, :])
        sampled_token = output_index_lookup[sampled_token_index]
        decoded_sentence += " " + sampled_token

        if sampled_token == "[eos]":
            break
    return decoded_sentence

def load_data(text_file):  
    with open(text_file) as f:
        lines = f.read().split("\n")[:-1]
    definition_pairs = []
    for line in lines:
        input, output = line.split("\t")
        output = "[bos] " + output + " [eos]"
        definition_pairs.append((input, output))
    return definition_pairs

def custom_standardization(input_string):
    lowercase = tf.strings.lower(input_string)
    return tf.strings.regex_replace(lowercase, "[%s]" % re.escape(strip_chars), "")

def format_dataset(input, output):
    input = input_vectorization(input)
    output = output_vectorization(output)
    return ({"encoder_inputs": input, "decoder_inputs": output[:, :-1],}, output[:, 1:])

def make_dataset(pairs):
    input_texts, output_texts = zip(*pairs)
    input_texts = list(input_texts)
    output_texts = list(output_texts)
    dataset = tf.data.Dataset.from_tensor_slices((input_texts, output_texts))
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(format_dataset)
    return dataset.shuffle(2048).prefetch(16).cache()

class TransformerEncoder(layers.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
        super(TransformerEncoder, self).__init__(**kwargs)
        self.embed_dim = embed_dim
        self.dense_dim = dense_dim
        self.num_heads = num_heads
        self.attention = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.dense_proj = keras.Sequential(
            [layers.Dense(dense_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.supports_masking = True

    def call(self, inputs, mask=None):
        if mask is not None:
            padding_mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], dtype="int32")
        attention_output = self.attention(
            query=inputs, value=inputs, key=inputs, attention_mask=padding_mask
        )
        proj_input = self.layernorm_1(inputs + attention_output)
        proj_output = self.dense_proj(proj_input)
        return self.layernorm_2(proj_input + proj_output)

class PositionalEmbedding(layers.Layer):
    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
        super(PositionalEmbedding, self).__init__(**kwargs)
        self.token_embeddings = layers.Embedding(
            input_dim=vocab_size, output_dim=embed_dim
        )
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=embed_dim
        )
        self.sequence_length = sequence_length
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim

    def call(self, inputs):
        length = tf.shape(inputs)[-1]
        positions = tf.range(start=0, limit=length, delta=1)
        embedded_tokens = self.token_embeddings(inputs)
        embedded_positions = self.position_embeddings(positions)
        return embedded_tokens + embedded_positions

    def compute_mask(self, inputs, mask=None):
        return tf.math.not_equal(inputs, 0)

class TransformerDecoder(layers.Layer):
    def __init__(self, embed_dim, latent_dim, num_heads, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.embed_dim = embed_dim
        self.latent_dim = latent_dim
        self.num_heads = num_heads
        self.attention_1 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.attention_2 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.dense_proj = keras.Sequential(
            [layers.Dense(latent_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.layernorm_3 = layers.LayerNormalization()
        self.supports_masking = True

    def call(self, inputs, encoder_outputs, mask=None):
        causal_mask = self.get_causal_attention_mask(inputs)
        if mask is not None:
            padding_mask = tf.cast(mask[:, tf.newaxis, :], dtype="int32")
            padding_mask = tf.minimum(padding_mask, causal_mask)

        attention_output_1 = self.attention_1(
            query=inputs, value=inputs, key=inputs, attention_mask=causal_mask
        )
        out_1 = self.layernorm_1(inputs + attention_output_1)

        attention_output_2 = self.attention_2(
            query=out_1,
            value=encoder_outputs,
            key=encoder_outputs,
            attention_mask=padding_mask,
        )
        out_2 = self.layernorm_2(out_1 + attention_output_2)

        proj_output = self.dense_proj(out_2)
        return self.layernorm_3(out_2 + proj_output)

    def get_causal_attention_mask(self, inputs):
        input_shape = tf.shape(inputs)
        batch_size, sequence_length = input_shape[0], input_shape[1]
        i = tf.range(sequence_length)[:, tf.newaxis]
        j = tf.range(sequence_length)
        mask = tf.cast(i >= j, dtype="int32")
        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
        mult = tf.concat(
            [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)],
            axis=0,
        )
        return tf.tile(mask, mult)


DATA_PATH_ABS = './Dataset1/Round'
epochs = 1  # We should increase it to at least 40 epochs, and check the corresponding accuracy


for fold in range(1,11):
    DATA_PATH = DATA_PATH_ABS + str(fold)
    print(f"FOLD {DATA_PATH}")
    training_file = join(DATA_PATH,'training_data.csv')
    training_pairs = load_data(training_file)

    test_file = join(DATA_PATH,'testing_data.csv')
    test_pairs = load_data(test_file)

    """We split the original training data to training and validation.
    """

    num_val_samples = int(0.1 * len(training_pairs))
    num_train_samples = len(training_pairs) - num_val_samples
    train_pairs = training_pairs[:num_train_samples]
    val_pairs = training_pairs[num_train_samples:]

    print(f"{len(train_pairs)} training pairs")
    print(f"{len(val_pairs)} validation pairs")
    print(f"{len(test_pairs)} test pairs")

    strip_chars = string.punctuation + "¿"
    strip_chars = strip_chars.replace("[", "")
    strip_chars = strip_chars.replace("]", "")

    vocab_size = 15000
    sequence_length = 20
    batch_size = 64

    input_vectorization = TextVectorization(
        max_tokens=vocab_size, output_mode="int", output_sequence_length=sequence_length,
    )
    output_vectorization = TextVectorization(
        max_tokens=vocab_size,
        output_mode="int",
        output_sequence_length=sequence_length + 1,
        standardize=custom_standardization,
    )
    train_input_texts = [pair[0] for pair in train_pairs]
    train_output_texts = [pair[1] for pair in train_pairs]
    input_vectorization.adapt(train_input_texts)
    output_vectorization.adapt(train_output_texts)

    train_ds = make_dataset(train_pairs)
    val_ds = make_dataset(val_pairs)




    for inputs, targets in train_ds.take(1):
        print(f'inputs["encoder_inputs"].shape: {inputs["encoder_inputs"].shape}')
        print(f'inputs["decoder_inputs"].shape: {inputs["decoder_inputs"].shape}')
        print(f"targets.shape: {targets.shape}")


    embed_dim = 256
    latent_dim = 2048
    num_heads = 8

    encoder_inputs = keras.Input(shape=(None,), dtype="int64", name="encoder_inputs")
    x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(encoder_inputs)
    encoder_outputs = TransformerEncoder(embed_dim, latent_dim, num_heads)(x)
    encoder = keras.Model(encoder_inputs, encoder_outputs)

    decoder_inputs = keras.Input(shape=(None,), dtype="int64", name="decoder_inputs")
    encoded_seq_inputs = keras.Input(shape=(None, embed_dim), name="decoder_state_inputs")
    x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(decoder_inputs)
    x = TransformerDecoder(embed_dim, latent_dim, num_heads)(x, encoded_seq_inputs)
    x = layers.Dropout(0.5)(x)
    decoder_outputs = layers.Dense(vocab_size, activation="softmax")(x)
    decoder = keras.Model([decoder_inputs, encoded_seq_inputs], decoder_outputs)

    decoder_outputs = decoder([decoder_inputs, encoder_outputs])
    transformer = keras.Model(
        [encoder_inputs, decoder_inputs], decoder_outputs, name="transformer"
    )
    startTime = datetime.now()
    start_time = startTime.strftime("%H:%M:%S")

    transformer.summary()
    transformer.compile(
        "rmsprop", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
    )
    transformer.fit(train_ds, epochs=epochs, validation_data=val_ds)

    #transformer.save_weights(join(DATA_PATH,'Transformer.h5'))

    transformer.save_weights(join(DATA_PATH,'Transformer220612.h5'))
    #transformer.load_weights(join(DATA_PATH,'Transformer.h5'))
    endTime = datetime.now()
    end_time = endTime.strftime("%H:%M:%S")
    print("Start TimeL =", start_time)
    print("End TimeL =", end_time)

    output_vocab = output_vectorization.get_vocabulary()
    output_index_lookup = dict(zip(range(len(output_vocab)), output_vocab))
    max_decoded_sentence_length = 20

    print("start testing here!")


    startTime = datetime.now()
    start_time = startTime.strftime("%H:%M:%S")

    test_input_texts = [pair[0] for pair in test_pairs]

    num_test = len(test_input_texts)
    result_file = join(DATA_PATH,'results.csv')

    with open(result_file, 'w+') as f1:
    	for seq_index in range(num_test):
        	# Take one sequence (part of the training set)
           	# for trying out decoding.
            input_seq = test_input_texts[seq_index : seq_index + 1]
            predicted_sequence = decode_sequence(input_seq)
            f1.write(test_input_texts[seq_index] + "\t" + predicted_sequence + "\n")
            print(test_input_texts[seq_index] + "\t" + predicted_sequence)
    endTime = datetime.now()
    end_time = endTime.strftime("%H:%M:%S")
    print("Start Time =", start_time)
    print("End Time =", end_time)