# ---
# 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
# ---

# + [markdown] papermill={"duration": 0.016151, "end_time": "2020-10-25T05:33:25.665485", "exception": false, "start_time": "2020-10-25T05:33:25.649334", "status": "completed"} tags=[]
# ## Parameter

# + papermill={"duration": 0.014613, "end_time": "2020-10-25T05:33:25.690669", "exception": false, "start_time": "2020-10-25T05:33:25.676056", "status": "completed"} tags=["parameters"]
# These parameters can be injected from Papermill
model_type = "pre_ln"
train_file = "wikitext-103-raw/wiki.train.raw"
valid_file = "wikitext-103-raw/wiki.valid.raw"
epochs = 10
batch_size = 2
max_learning_rate = 1e-4
warmup_steps = 0
save_model_dir = "tfchat_model"
clipnorm = 1.0
fp16 = False
# -

# Assert parameters
assert model_type in ["pre_ln", "post_ln", "min_gpt", "transformers"]

# ## Installation

# !apt install -y git
# !pip install git+https://github.com/noriyukipy/tfchat@ad516f2

# + [markdown] papermill={"duration": 0.010568, "end_time": "2020-10-25T05:33:25.736850", "exception": false, "start_time": "2020-10-25T05:33:25.726282", "status": "completed"} tags=[]
# ## Configure GPU

# + papermill={"duration": 0.920123, "end_time": "2020-10-25T05:33:26.667484", "exception": false, "start_time": "2020-10-25T05:33:25.747361", "status": "completed"} tags=[]
from tfchat.utils import set_memory_growth
from tfchat.utils import set_mixed_precision_policy

# + papermill={"duration": 0.55073, "end_time": "2020-10-25T05:33:27.228969", "exception": false, "start_time": "2020-10-25T05:33:26.678239", "status": "completed"} tags=[]
set_memory_growth()
# -

if fp16:
    set_mixed_precision_policy()

# + [markdown] papermill={"duration": 0.010916, "end_time": "2020-10-25T05:33:27.251283", "exception": false, "start_time": "2020-10-25T05:33:27.240367", "status": "completed"} tags=[]
# ## Setup tokenizer

# + papermill={"duration": 1.424651, "end_time": "2020-10-25T05:33:28.686639", "exception": false, "start_time": "2020-10-25T05:33:27.261988", "status": "completed"} tags=[]
# Install transformers by HuggingFace to use GPT2 tokenizer
# ! pip install transformers==3.4.0
# Enable widgetsnbextention to avoid the following error when running GPT2.from_pretrained method
#     ImportError: IProgress not found. Please update jupyter and ipywidgets.
# ! jupyter nbextension enable --py widgetsnbextension

# + papermill={"duration": 1.701922, "end_time": "2020-10-25T05:33:30.435314", "exception": false, "start_time": "2020-10-25T05:33:28.733392", "status": "completed"} tags=[]
# setup tokenizer
from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# + [markdown] papermill={"duration": 0.012068, "end_time": "2020-10-25T05:33:30.459683", "exception": false, "start_time": "2020-10-25T05:33:30.447615", "status": "completed"} tags=[]
# ## Prepare model config

# + papermill={"duration": 0.084744, "end_time": "2020-10-25T05:33:30.556263", "exception": false, "start_time": "2020-10-25T05:33:30.471519", "status": "completed"} tags=[]
from tfchat.configs import GPT2SmallConfig

config = GPT2SmallConfig()

# Set the larger number of vocab size than 33,278, which is the vocab size of Wikitext-2
config.vocab_size = tokenizer.vocab_size

# + papermill={"duration": 0.020325, "end_time": "2020-10-25T05:33:30.588878", "exception": false, "start_time": "2020-10-25T05:33:30.568553", "status": "completed"} tags=[]
config

# + [markdown] papermill={"duration": 0.012416, "end_time": "2020-10-25T05:33:30.613934", "exception": false, "start_time": "2020-10-25T05:33:30.601518", "status": "completed"} tags=[]
# ## Prepare Dataset

# + papermill={"duration": 0.017599, "end_time": "2020-10-25T05:33:30.643992", "exception": false, "start_time": "2020-10-25T05:33:30.626393", "status": "completed"} tags=[]
from pathlib import Path
from urllib.request import urlretrieve
import zipfile
import numpy as np


def encode_file(_tokenizer, _filepath):
    ids = []
    with open(_filepath) as f:
        for line in f.readlines():
            text = line.strip("\n")
            ids.extend(_tokenizer.encode(text))

    return np.array(ids, dtype=np.int32)


# + papermill={"duration": 12.24382, "end_time": "2020-10-25T05:33:42.900304", "exception": false, "start_time": "2020-10-25T05:33:30.656484", "status": "completed"} tags=[]
train_ids = encode_file(tokenizer, train_file)
valid_ids = encode_file(tokenizer, valid_file)

# + papermill={"duration": 0.017183, "end_time": "2020-10-25T05:33:42.930359", "exception": false, "start_time": "2020-10-25T05:33:42.913176", "status": "completed"} tags=[]
print("Train:", train_ids.shape)
print("Valid:", valid_ids.shape)

# + papermill={"duration": 0.016829, "end_time": "2020-10-25T05:33:42.959985", "exception": false, "start_time": "2020-10-25T05:33:42.943156", "status": "completed"} tags=[]
print(train_ids.shape)
print(valid_ids.shape)

# + papermill={"duration": 0.484628, "end_time": "2020-10-25T05:33:43.457723", "exception": false, "start_time": "2020-10-25T05:33:42.973095", "status": "completed"} tags=[]
from tfchat.data import BlockDataset


dataset = BlockDataset(block_size=config.context_size, batch_size=batch_size)

train_dataset = dataset.build(train_ids, shuffle=True)
valid_dataset = dataset.build(valid_ids, shuffle=False)

# + papermill={"duration": 0.017, "end_time": "2020-10-25T05:33:43.488196", "exception": false, "start_time": "2020-10-25T05:33:43.471196", "status": "completed"} tags=[]
num_train_steps = len([_ for _ in train_dataset])
num_valid_steps = len([_ for _ in valid_dataset])
print("Train steps:", num_train_steps)
print("Valid steps:", num_valid_steps)
# -

# ## Transformers model implementation

from transformers import TFGPT2LMHeadModel
from transformers import GPT2Config
import tensorflow.keras as keras
import tensorflow as tf
from tfchat.models import create_combined_mask


# +
class TransformersGPT2(keras.Model):
    def __init__(self, config):
        super().__init__()
        tf_config = GPT2Config(
            n_layers=config.num_layers,
            n_embd=config.d_model,
            n_head=config.num_heads,
            n_inner=config.d_ff,
            vocab_size=config.vocab_size,
            n_ctx=config.context_size,
            n_positions=config.context_size,
            attn_pdrop=config.attention_dropout_rate,
            resid_pdrop=config.residual_dropout_rate,
            embd_pdrop=config.embedding_dropout_rate,
            layer_norm_epsilon=config.epsilon,
            activation_function="gelu_new",  # Default value of transformers implementation
            
        )
        self._decoder = TFGPT2LMHeadModel(tf_config)
        
    def call(self, inputs, training):
        inputs = tf.cast(inputs, tf.int32)
        x = self._decoder(inputs, training=training)
        return x[0]



# + [markdown] papermill={"duration": 0.013209, "end_time": "2020-10-25T05:33:43.514599", "exception": false, "start_time": "2020-10-25T05:33:43.501390", "status": "completed"} tags=[]
# ## Prepare Model

# + papermill={"duration": 0.019812, "end_time": "2020-10-25T05:33:43.547540", "exception": false, "start_time": "2020-10-25T05:33:43.527728", "status": "completed"} tags=[]
from tfchat.losses import PaddingLoss
from tfchat.schedules import WarmupLinearDecay
import tensorflow.keras as keras



def train(_model, _train_dataset, _valid_dataset, _epochs, _warmup_steps, _num_train_steps, _max_learning_rate, _clipnorm):
    schedule = WarmupLinearDecay(max_learning_rate=_max_learning_rate,
                                 warmup_steps=_warmup_steps,
                                 training_steps=_num_train_steps*_epochs)
    optimizer = keras.optimizers.Adam(schedule, beta_1=0.9, beta_2=0.999, epsilon=1e-8, clipnorm=_clipnorm)
    _model.compile(loss=PaddingLoss(), optimizer=optimizer)


    history = _model.fit(
        _train_dataset,
        validation_data=_valid_dataset,
        epochs=_epochs,
        callbacks=[
            keras.callbacks.EarlyStopping(patience=1, restore_best_weights=True),
            # If you want to save chekcpoints, remove the next comment out
            #keras.callbacks.ModelCheckpoint("keras_model/", save_best_only=True)
        ],
        verbose=2,
    )

# -

if model_type == "pre_ln":
    from tfchat.models import PreLNDecoder
    model = PreLNDecoder(config)
elif model_type == "post_ln":
    from tfchat.models import PostLNDecoder 
    model = PostLNDecoder(config)
elif model_type == "transformers":
    model = TransformersGPT2(config)
elif model_type == "min_gpt":
    from mingpt.model import GPT, GPTConfig
    mconf = GPTConfig(config.vocab_size, config.context_size,
                      n_layer=config.num_layers, n_head=config.num_heads, n_embd=config.d_model)
    model = GPT(mconf)
else:
    raise Exception("Model type is wrong")

model.build(input_shape=(None, config.context_size))
model.summary()

# + papermill={"duration": 535.285785, "end_time": "2020-10-25T05:42:39.126238", "exception": false, "start_time": "2020-10-25T05:33:43.840453", "status": "completed"} tags=[]
train(model, train_dataset, valid_dataset, epochs, warmup_steps, num_train_steps, max_learning_rate, clipnorm)

# +
from tfchat.eval import perplexity

print("Validation PPL:", perplexity(model, valid_dataset))

# +
from tfchat.utils import save_model

save_model(save_model_dir, model, config)
# -



StarcoderdataJupyterScriptsDedupFiltered