# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# -

from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.layers import common_attention
from tensor2tensor.utils import registry
from tensor2tensor import problems
import tensorflow as tf
import os
import logging
import sentencepiece as spm
import transformer_tag
from tensor2tensor.layers import modalities

# +
vocab = 'sp10m.cased.t5.model'
sp = spm.SentencePieceProcessor()
sp.Load(vocab)

class Encoder:
    def __init__(self, sp):
        self.sp = sp
        self.vocab_size = sp.GetPieceSize() + 100

    def encode(self, s):
        return self.sp.EncodeAsIds(s)

    def decode(self, ids, strip_extraneous = False):
        return self.sp.DecodeIds(list(ids))


# +
d = [
    {'class': 0, 'Description': 'PAD', 'salah': '', 'betul': ''},
    {
        'class': 1,
        'Description': 'kesambungan subwords',
        'salah': '',
        'betul': '',
    },
    {
        'class': 2,
        'Description': 'tiada kesalahan',
        'salah': '',
        'betul': '',
    },
    {
        'class': 3,
        'Description': 'kesalahan frasa nama, Perkara yang diterangkan mesti mendahului "penerang"',
        'salah': 'Cili sos',
        'betul': 'sos cili',
    },
    {
        'class': 4,
        'Description': 'kesalahan kata jamak',
        'salah': 'mereka-mereka',
        'betul': 'mereka',
    },
    {
        'class': 5,
        'Description': 'kesalahan kata penguat',
        'salah': 'sangat tinggi sekali',
        'betul': 'sangat tinggi',
    },
    {
        'class': 6,
        'Description': 'kata adjektif dan imbuhan "ter" tanpa penguat.',
        'salah': 'Sani mendapat markah yang tertinggi sekali.',
        'betul': 'Sani mendapat markah yang tertinggi.',
    },
    {
        'class': 7,
        'Description': 'kesalahan kata hubung',
        'salah': 'Sally sedang membaca bila saya tiba di rumahnya.',
        'betul': 'Sally sedang membaca apabila saya tiba di rumahnya.',
    },
    {
        'class': 8,
        'Description': 'kesalahan kata bilangan',
        'salah': 'Beribu peniaga tidak membayar cukai pendapatan.',
        'betul': 'Beribu-ribu peniaga tidak membayar cukai pendapatan',
    },
    {
        'class': 9,
        'Description': 'kesalahan kata sendi',
        'salah': 'Umar telah berpindah daripada sekolah ini bulan lalu.',
        'betul': 'Umar telah berpindah dari sekolah ini bulan lalu.',
    },
    {
        'class': 10,
        'Description': 'kesalahan penjodoh bilangan',
        'salah': 'Setiap orang pelajar',
        'betul': 'Setiap pelajar.',
    },
    {
        'class': 11,
        'Description': 'kesalahan kata ganti diri',
        'salah': 'Pencuri itu telah ditangkap. Beliau dibawa ke balai polis.',
        'betul': 'Pencuri itu telah ditangkap. Dia dibawa ke balai polis.',
    },
    {
        'class': 12,
        'Description': 'kesalahan ayat pasif',
        'salah': 'Cerpen itu telah dikarang oleh saya.',
        'betul': 'Cerpen itu telah saya karang.',
    },
    {
        'class': 13,
        'Description': 'kesalahan kata tanya',
        'salah': 'Kamu berasal dari manakah ?',
        'betul': 'Kamu berasal dari mana ?',
    },
    {
        'class': 14,
        'Description': 'kesalahan tanda baca',
        'salah': 'Kamu berasal dari manakah .',
        'betul': 'Kamu berasal dari mana ?',
    },
    {
        'class': 15,
        'Description': 'kesalahan kata kerja tak transitif',
        'salah': 'Dia kata kepada saya',
        'betul': 'Dia berkata kepada saya',
    },
    {
        'class': 16,
        'Description': 'kesalahan kata kerja transitif',
        'salah': 'Dia suka baca buku',
        'betul': 'Dia suka membaca buku',
    },
    {
        'class': 17,
        'Description': 'penggunaan kata yang tidak tepat',
        'salah': 'Tembuk Besar negeri Cina dibina oleh Shih Huang Ti.',
        'betul': 'Tembok Besar negeri Cina dibina oleh Shih Huang Ti',
    },
]


class Tatabahasa:
    def __init__(self, d):
        self.d = d
        self.kesalahan = {i['Description']: no for no, i in enumerate(self.d)}
        self.reverse_kesalahan = {v: k for k, v in self.kesalahan.items()}
        self.vocab_size = len(self.d)

    def encode(self, s):
        return [self.kesalahan[i] for i in s]

    def decode(self, ids, strip_extraneous = False):
        return [self.reverse_kesalahan[i] for i in ids]


# -

@registry.register_problem
class Grammar(text_problems.Text2TextProblem):
    """grammatical error correction."""

    def feature_encoders(self, data_dir):
        encoder = Encoder(sp)
        t = Tatabahasa(d)
        return {'inputs': encoder, 'targets': encoder, 'targets_error_tag': t}

    def hparams(self, defaults, model_hparams):
        super(Grammar, self).hparams(defaults, model_hparams)
        if 'use_error_tags' not in model_hparams:
            model_hparams.add_hparam('use_error_tags', True)
        if 'middle_prediction' not in model_hparams:
            model_hparams.add_hparam('middle_prediction', False)
        if 'middle_prediction_layer_factor' not in model_hparams:
            model_hparams.add_hparam('middle_prediction_layer_factor', 2)
        if 'ffn_in_prediction_cascade' not in model_hparams:
            model_hparams.add_hparam('ffn_in_prediction_cascade', 1)
        if 'error_tag_embed_size' not in model_hparams:
            model_hparams.add_hparam('error_tag_embed_size', 12)
        if model_hparams.use_error_tags:
            defaults.modality[
                'targets_error_tag'
            ] = modalities.ModalityType.SYMBOL
            error_tag_vocab_size = self._encoders[
                'targets_error_tag'
            ].vocab_size
            defaults.vocab_size['targets_error_tag'] = error_tag_vocab_size

    def example_reading_spec(self):
        data_fields, _ = super(Grammar, self).example_reading_spec()
        data_fields['targets_error_tag'] = tf.VarLenFeature(tf.int64)
        return data_fields, None

    @property
    def approx_vocab_size(self):
        return 32100

    @property
    def is_generate_per_split(self):
        return False

    @property
    def dataset_splits(self):
        return [
            {'split': problem.DatasetSplit.TRAIN, 'shards': 200},
            {'split': problem.DatasetSplit.EVAL, 'shards': 1},
        ]


DATA_DIR = os.path.expanduser('t2t-tatabahasa/data')
TMP_DIR = os.path.expanduser('t2t-tatabahasa/tmp')
TRAIN_DIR = os.path.expanduser('t2t-tatabahasa/train-small')

PROBLEM = 'grammar'
t2t_problem = problems.problem(PROBLEM)

MODEL = 'transformer_tag'
HPARAMS = 'transformer_base'

from tensor2tensor.utils.trainer_lib import create_run_config, create_experiment
from tensor2tensor.utils.trainer_lib import create_hparams
from tensor2tensor.utils import registry
from tensor2tensor import models
from tensor2tensor import problems
from tensor2tensor.utils import trainer_lib

# +
X = tf.placeholder(tf.int32, [None, None], name = 'x_placeholder')
Y = tf.placeholder(tf.int32, [None, None], name = 'y_placeholder')
targets_error_tag = tf.placeholder(tf.int32, [None, None], 'error_placeholder')
X_seq_len = tf.count_nonzero(X, 1, dtype=tf.int32)
maxlen_decode = tf.reduce_max(X_seq_len)

x = tf.expand_dims(tf.expand_dims(X, -1), -1)
y = tf.expand_dims(tf.expand_dims(Y, -1), -1)
targets_error_tag_ = tf.expand_dims(tf.expand_dims(targets_error_tag, -1), -1)

features = {
    "inputs": x,
    "targets": y,
    "target_space_id": tf.constant(1, dtype=tf.int32),
    'targets_error_tag': targets_error_tag,
}
Modes = tf.estimator.ModeKeys
hparams = trainer_lib.create_hparams(HPARAMS, data_dir=DATA_DIR, problem_name=PROBLEM)

# +
hparams.filter_size = 2048
hparams.hidden_size = 512
hparams.num_heads = 8
hparams.num_hidden_layers = 6
hparams.vocab_divisor = 128
hparams.dropout = 0.1
hparams.max_length = 256

# LM
hparams.label_smoothing = 0.0
hparams.shared_embedding_and_softmax_weights = False
hparams.eval_drop_long_sequences = True
hparams.max_length = 256
hparams.multiproblem_mixing_schedule = 'pretrain'

# tpu
hparams.symbol_modality_num_shards = 1
hparams.attention_dropout_broadcast_dims = '0,1'
hparams.relu_dropout_broadcast_dims = '1'
hparams.layer_prepostprocess_dropout_broadcast_dims = '1'
# -

model = registry.model(MODEL)(hparams, Modes.PREDICT)

# +
# logits = model(features)
# logits

# sess = tf.InteractiveSession()
# sess.run(tf.global_variables_initializer())
# l = sess.run(logits, feed_dict = {X: [[10,10, 10, 10,10,1],[10,10, 10, 10,10,1]],
#                              Y: [[10,10, 10, 10,10,1],[10,10, 10, 10,10,1]],
#                              targets_error_tag: [[10,10, 10, 10,10,1],
#                                                 [10,10, 10, 10,10,1]]})

# +
features = {
    "inputs": x,
    "target_space_id": tf.constant(1, dtype=tf.int32),
}

with tf.variable_scope(tf.get_variable_scope(), reuse = False):
    fast_result = model._greedy_infer(features, maxlen_decode)
# -

result_seq = tf.identity(fast_result['outputs'], name = 'greedy')
result_tag = tf.identity(fast_result['outputs_tag'], name = 'tag_greedy')

# +
from tensor2tensor.layers import common_layers

def accuracy_per_sequence(predictions, targets, weights_fn = common_layers.weights_nonzero):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(predictions, targets)
    weights = weights_fn(padded_labels)
    padded_labels = tf.to_int32(padded_labels)
    padded_predictions = tf.to_int32(padded_predictions)
    not_correct = tf.to_float(tf.not_equal(padded_predictions, padded_labels)) * weights
    axis = list(range(1, len(padded_predictions.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return tf.reduce_mean(correct_seq)

def padded_accuracy(predictions, targets, weights_fn = common_layers.weights_nonzero):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(predictions, targets)
    weights = weights_fn(padded_labels)
    padded_labels = tf.to_int32(padded_labels)
    padded_predictions = tf.to_int32(padded_predictions)
    n = tf.to_float(tf.equal(padded_predictions, padded_labels)) * weights
    d = tf.reduce_sum(weights)
    return tf.reduce_sum(n) / d


# -

acc_seq = padded_accuracy(result_seq, Y)
acc_tag = padded_accuracy(result_tag, targets_error_tag)

ckpt_path = tf.train.latest_checkpoint(os.path.join(TRAIN_DIR))
ckpt_path

sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())

var_lists = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
saver = tf.train.Saver(var_list = var_lists)
saver.restore(sess, ckpt_path)

# +
import pickle

with open('../pure-text/dataset-tatabahasa.pkl', 'rb') as fopen:
    data = pickle.load(fopen)

encoder = Encoder(sp)


# -

def get_xy(row, encoder):
    x, y, tag = [], [], []

    for i in range(len(row[0])):
        t = encoder.encode(row[0][i][0])
        y.extend(t)
        t = encoder.encode(row[1][i][0])
        x.extend(t)
        tag.extend([row[1][i][1]] * len(t))

    # EOS
    x.append(1)
    y.append(1)
    tag.append(0)

    return x, y, tag


import numpy as np

x, y, tag = get_xy(data[10], encoder)

e = encoder.encode('Pilih mana jurusan yang sesuai dengan kebolehan anda dalam peperiksaan Sijil Pelajaran Malaysia semasa memohon kemasukan ke institusi pengajian tinggi.') + [1]

r = sess.run(fast_result, 
         feed_dict = {X: [e]})

r['outputs_tag']

encoder.decode(r['outputs'][0].tolist())

encoder.decode(x)

encoder.decode(y)

hparams.problem.example_reading_spec()[0]


def parse(serialized_example):

    data_fields = hparams.problem.example_reading_spec()[0]
    features = tf.parse_single_example(
        serialized_example, features = data_fields
    )
    for k in features.keys():
        features[k] = features[k].values

    return features


dataset = tf.data.TFRecordDataset('t2t-tatabahasa/data/grammar-dev-00000-of-00001')
dataset = dataset.map(parse, num_parallel_calls=32)
dataset = dataset.padded_batch(32, 
    padded_shapes = {
    'inputs': tf.TensorShape([None]),
    'targets': tf.TensorShape([None]),
    'targets_error_tag': tf.TensorShape([None])
    },
    padding_values = {
        'inputs': tf.constant(0, dtype = tf.int64),
        'targets': tf.constant(0, dtype = tf.int64),
        'targets_error_tag': tf.constant(0, dtype = tf.int64),
    })
dataset = dataset.make_one_shot_iterator().get_next()
dataset

seqs, tags = [], []
index = 0
while True:
    try:
        d = sess.run(dataset)
        s, t = sess.run([acc_seq, acc_tag], feed_dict = {X:d['inputs'], 
                                              Y: d['targets'], 
                                              targets_error_tag: d['targets_error_tag']})
        seqs.append(s)
        tags.append(t)
        print(f'done {index}')
        index += 1
    except:
        break

np.mean(seqs), np.mean(tags)

saver = tf.train.Saver(tf.trainable_variables())
saver.save(sess, 'transformertag-small/model.ckpt')

strings = ','.join(
    [
        n.name
        for n in tf.get_default_graph().as_graph_def().node
        if ('Variable' in n.op
        or 'Placeholder' in n.name
        or 'greedy' in n.name
        or 'tag_greedy' in n.name
        or 'x_placeholder' in n.name
        or 'self/Softmax' in n.name)
        and 'adam' not in n.name
        and 'beta' not in n.name
        and 'global_step' not in n.name
        and 'modality' not in n.name
        and 'Assign' not in n.name
    ]
)
strings.split(',')


def freeze_graph(model_dir, output_node_names):

    if not tf.gfile.Exists(model_dir):
        raise AssertionError(
            "Export directory doesn't exists. Please specify an export "
            'directory: %s' % model_dir
        )

    checkpoint = tf.train.get_checkpoint_state(model_dir)
    input_checkpoint = checkpoint.model_checkpoint_path

    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])
    output_graph = absolute_model_dir + '/frozen_model.pb'
    clear_devices = True
    with tf.Session(graph = tf.Graph()) as sess:
        saver = tf.train.import_meta_graph(
            input_checkpoint + '.meta', clear_devices = clear_devices
        )
        saver.restore(sess, input_checkpoint)
        output_graph_def = tf.graph_util.convert_variables_to_constants(
            sess,
            tf.get_default_graph().as_graph_def(),
            output_node_names.split(','),
        )
        with tf.gfile.GFile(output_graph, 'wb') as f:
            f.write(output_graph_def.SerializeToString())
        print('%d ops in the final graph.' % len(output_graph_def.node))


freeze_graph('transformertag-small', strings)


def load_graph(frozen_graph_filename):
    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def)
    return graph


g = load_graph('transformertag-small/frozen_model.pb')
x = g.get_tensor_by_name('import/x_placeholder:0')
greedy = g.get_tensor_by_name('import/greedy:0')
tag_greedy = g.get_tensor_by_name('import/tag_greedy:0')
test_sess = tf.InteractiveSession(graph = g)

test_sess.run([greedy, tag_greedy], feed_dict = {x:d['inputs']})

import tensorflow as tf
from tensorflow.tools.graph_transforms import TransformGraph
from glob import glob
tf.set_random_seed(0)

import tensorflow_text
import tf_sentencepiece

# +
transforms = ['add_default_attributes',
             'remove_nodes(op=Identity, op=CheckNumerics, op=Dropout)',
             'fold_constants(ignore_errors=true)',
             'fold_batch_norms',
             'fold_old_batch_norms',
             'quantize_weights(fallback_min=-10, fallback_max=10)',
             'strip_unused_nodes',
             'sort_by_execution_order']

pb = 'transformertag-small/frozen_model.pb'
input_graph_def = tf.GraphDef()
with tf.gfile.FastGFile(pb, 'rb') as f:
    input_graph_def.ParseFromString(f.read())
    
transformed_graph_def = TransformGraph(input_graph_def, 
                                       ['x_placeholder'],
                                       ['greedy', 'tag_greedy'], transforms)

with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:
    f.write(transformed_graph_def.SerializeToString())
# -

g = load_graph('transformertag-small/frozen_model.pb.quantized')
x = g.get_tensor_by_name('import/x_placeholder:0')
greedy = g.get_tensor_by_name('import/greedy:0')
tag_greedy = g.get_tensor_by_name('import/tag_greedy:0')
test_sess = tf.InteractiveSession(graph = g)

test_sess.run([greedy, tag_greedy], feed_dict = {x:d['inputs']})

StarcoderdataJupyterScriptsDedupFiltered