# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3.8.5 64-bit ('usr')
#     name: python3
# ---

# + id="fTZYaL1ANwxu"
# Following https://github.com/huggingface/notebooks/blob/master/examples/translation.ipynb

# + id="566GHCfaILnG"
# Load the TensorBoard notebook extension
# %load_ext tensorboard
# Clear any logs from previous runs
# !rm -rf ./logs/ 

# + colab={"base_uri": "https://localhost:8080/"} id="CpHoG6l5Nx1E" executionInfo={"status": "ok", "timestamp": 1628500934784, "user_tz": -120, "elapsed": 97113, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="b86e5bee-1850-4c55-d78a-104ecfc4e8d1"
from google.colab import drive
drive.mount('drive')

# + colab={"base_uri": "https://localhost:8080/"} id="FUM2HkRlKpUE" executionInfo={"status": "ok", "timestamp": 1628501035852, "user_tz": -120, "elapsed": 524, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="3fa0fd22-8c50-4f14-e19d-4136ae7e264e"
# !unzip lc-quad-wikidata-2021-07-09.zip

# + id="EXgnU7gQOGGS" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628500934787, "user_tz": -120, "elapsed": 18, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="b8068c78-67ab-4f53-a8c7-404b6d0f5b89"
# !nvidia-smi

# + id="Q9TjriZbNwxy"
model_checkpoint = 't5-small'
gdir = 'drive/My Drive/Colab Notebooks/t5-2021-07-11/'
model_name='sparql-translator-t5-2021-08-08-large'
model_path='./models/'+model_name
ds_path= 'lc-quad-wikidata-2021-07-09'

# + id="Acx8YqC8Ok-H"
# !mkdir models

# + id="fz6l2bnROu-W"
# !pip install datasets transformers sacrebleu -qqq

# + id="WGFsMmwlNwxz"
from datasets import load_dataset, load_metric, Dataset, load_from_disk
raw_datasets = load_from_disk(ds_path)

# + id="ak3Fc0pDUG8_" colab={"base_uri": "https://localhost:8080/", "height": 253} executionInfo={"status": "ok", "timestamp": 1628501050514, "user_tz": -120, "elapsed": 4102, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="cbad26bd-65b0-4831-b092-51b7f930e388"
# !pip install tqdm==4.49.0

# + id="-dRUZHjeNwxz"
from tqdm import tqdm

# + id="gqgySQUJFZev" colab={"base_uri": "https://localhost:8080/", "height": 49, "referenced_widgets": ["ed506f08601040e8b373290d1122f473", "75cb744503cb4d59aa8464af590b8912", "08860bcdbba644188c36feda590e42da", "9920ab5ff73445a6aaa8f4510626a039", "531a28ce7f014769ad297b0c9c71635b", "63752066a91c469bb17cb3b0c9822f53", "1afe950c10bb40179af4ae10212b769f", "<KEY>", "67c4c17e413b4d0fa035bb892abf240f", "<KEY>", "<KEY>"]} executionInfo={"status": "ok", "timestamp": 1628501051236, "user_tz": -120, "elapsed": 739, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="9e79ef5e-aa1f-4b2b-c23d-6cd68cd2b720"
metric = load_metric("sacrebleu")

# + id="HGDdBpGfNwxz"
# Preprocessing

# + id="UnY4aKkPNwx0" colab={"base_uri": "https://localhost:8080/", "height": 113, "referenced_widgets": ["c539809a354b4d8b91fc0962f8897225", "629731748a8b4158ab7ab9ed5d01d1dd", "691f8e5dc8224226bce91db3a0d4f0de", "ff9cdc17850a4e6183d6f7b2c31bd16a", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "eb7e8ddc9217450d9f8ad8b9c8db82af", "63f778a0af934159a50c760596e80f3f", "d0218ea78c9a4f238e83e35145b3d719", "c7073fe66c864ca0b47e6e146a1cba8e", "06f309eaffcb49c1a5aa792bed9b348c", "<KEY>", "<KEY>", "01699bed9a9e45edae815b20fe0163a7", "bf0501456ece4a7d99ecda20d80586a3", "<KEY>", "<KEY>", "<KEY>", "370111135ec5432c85a8662668279c8a", "f983e734d0f940098f1feb550f6dd7bc", "514e3b2e5cd446808a1899e0b49c16e2", "<KEY>", "3338874a7d8c4e24864f7c8a79b6b494", "b136deab206a440ca0a8954559975622", "44cc2e9972584bea94d861fed1b63aaf", "8588827b5a6e4087be1345210999c9e3", "8b398d8b1d8d459fb0d441c7bec3ac62", "bb9e98a3d29c499f85e4ab3323f0ff30", "b67b13dad4e14233a46ba75f63ced5d1", "bff85e6eb4be444c8b80ce7a450fe90e", "cb044677cf67427985bb4ef576e76db6"]} executionInfo={"status": "ok", "timestamp": 1628501062561, "user_tz": -120, "elapsed": 11330, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="695e7074-cc33-469e-a323-931bfe4f173a"
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

# + id="LtJhGYmEaJAr" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501062563, "user_tz": -120, "elapsed": 19, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="534fc079-ec03-4df3-b97e-63ea8fd58ab0"
print(raw_datasets['test']['translation'][0]['sparql'])

# + id="5iL-TpK3mUk8" colab={"base_uri": "https://localhost:8080/", "height": 35} executionInfo={"status": "ok", "timestamp": 1628501062566, "user_tz": -120, "elapsed": 16, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="39c012ac-3d8d-462b-bd87-6d546f96fe1d"
tokenizer.decode(tokenizer(raw_datasets['test']['translation'][0]['sparql'])['input_ids'])

# + id="EZdjRilTNwx1" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501062568, "user_tz": -120, "elapsed": 17, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="263cbad6-1635-4953-a26e-89652d18a3c2"
print(raw_datasets['test']['translation'][0]['en'],tokenizer(raw_datasets['test']['translation'][0]['en']))

# + id="tNf9AWfpNwx1"
prefix = "translate English to Sparql: "

# + tags=[] id="kkidcZ_hNwx2" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501062954, "user_tz": -120, "elapsed": 398, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="a9b17439-f19a-4e02-d507-fa80f5d40b6e"
max_input_length = 0 
max_target_length = 0
for d in tqdm(raw_datasets['train']['translation']):
    len_en = len(d['en'])
    len_qry = len(d['sparql'])
    if len_en > max_input_length: max_input_length=len_en
    if len_qry > max_target_length: max_target_length=len_qry

# + id="JyEtHVN6Nwx2" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501062955, "user_tz": -120, "elapsed": 11, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="f1b94bf9-87e6-4964-d85e-9965252eefe3"
print(max_input_length, max_target_length)

# + id="OXiURhR-Nwx3"
source_lang = "en"
target_lang = "sparql"

def preprocess_function(examples):
    inputs = []
    targets= []
    for ex in examples["translation"]:
      inputs.append(prefix + ex[source_lang])
      targets.append(ex[target_lang])
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

    # Setup the tokenizer for targets
    with tokenizer.as_target_tokenizer():
      labels = tokenizer(targets, max_length=max_target_length, truncation=True)
      
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs


# + id="qmaK1UX0Nwx3" colab={"base_uri": "https://localhost:8080/", "height": 81, "referenced_widgets": ["df458fbd27154e9da34ad7c40dac4972", "c761a0828bbe452495badcf480f569fa", "e448844fbe91486f9fcc3fbe0a8335f9", "fde4d7035795479a887f055cd08b023a", "<KEY>", "3ef1ace60831437fb2513aa76f5d13c7", "<KEY>", "f43fe3695925442e93b6348796ce9f4f", "607796fb886946e68a2918d30e462295", "<KEY>", "<KEY>", "<KEY>", "2678da2eb22e496db3429559ad05e5a0", "7ef1c54a35b54cdea4ce6bbe80fc4ee6", "abdfd23f71f647899182974013ada503", "965e825363f34036b3b1b95ecae66b32", "7c0f763aadbe4403ac8b8b1ab69a7155", "8fdfb3e4c13a4032b2b342866a060c22", "<KEY>", "<KEY>", "a0fb2219fb7d4f0ca13ba9e684c669aa", "347f90c961784ee7a66b798e8e07c043"]} executionInfo={"status": "ok", "timestamp": 1628501072550, "user_tz": -120, "elapsed": 9602, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="1a4cd3bc-0e34-4d36-ec8a-ad274efba063"
tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)

# + id="knRoh_8wNwx4" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501072552, "user_tz": -120, "elapsed": 14, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="9a038412-9bf0-4ccb-9d01-63b2ee215d3f"
tokenized_datasets

# + id="eUiaZ1Y0Nwx4"
# Fine-tuning the model

# + id="rUJNfPjvNwx5" colab={"base_uri": "https://localhost:8080/", "height": 49, "referenced_widgets": ["6f260fef1c70436d9549801c615d3bb3", "92625a6d0c2144f7b728f682917aeb55", "c8c3ee1999df430f88a06d594e6b3d4e", "fec6412bc353490687c6655fb361f109", "41e671a7afe942d7955a7fc2aac9dc05", "b7a5e60e91be432f90189ae567156191", "329983b598f44153ab1c38490e83d4a5", "abe8706090ce4645add8d60736eb2f31", "cc0240d6105a419e92c9d063e2aebc78", "46f24cb47059491a8db4f8fbcffe8aff", "42a73d2ebd9246f08233dfbe85687608"]} executionInfo={"status": "ok", "timestamp": 1628501083180, "user_tz": -120, "elapsed": 10636, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="8f3c8f5d-46ae-421a-b45f-9f403a99ff3a"
from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

# + id="Qt8rECNwNwx5"
batch_size = 8
args = Seq2SeqTrainingArguments(
    model_name,
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=1,
    predict_with_generate=True,
    fp16=True,
)

# + id="wHZPXO3lNwx5"
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

# + id="UO8h_ilTLM3Y" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501083181, "user_tz": -120, "elapsed": 18, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="acfa5a4c-fee2-4335-910e-87892ed50fc7"
fake_preds = ["hello there", "general kenobi"]
fake_labels = [["hello there"], ["general kenobi"]]
metric.compute(predictions=fake_preds, references=fake_labels)

# + id="7TzcEj3jKSlD"
# The last thing to define for our Seq2SeqTrainer is how to compute 
# the metrics from the predictions. We need to define 
# a function for this, which will just use the metric we loaded earlier, 
# and we have to do a bit of pre-processing to decode the predictions into texts:

import numpy as np

def postprocess_text(preds, labels):
    preds = [pred.replace('?',' ?').replace('.', ' .').strip() for pred in preds]
    labels = [[label.replace('?',' ?').replace('.', ' .').strip()] for label in labels]

    return preds, labels

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # Some simple post-processing
    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels)
    result = {"bleu": result["score"]}

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    result = {k: round(v, 4) for k, v in result.items()}
    return result


# + id="9ZhWPx09Nwx6" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628501093036, "user_tz": -120, "elapsed": 9869, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="6daf2183-1ec8-4cd3-88e1-7ede335987bd"
trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

# + colab={"base_uri": "https://localhost:8080/"} id="SgIrXB8xNTld" executionInfo={"status": "ok", "timestamp": 1628501355186, "user_tz": -120, "elapsed": 262167, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="34a7da8b-0e20-42b2-c82d-9549260868eb"
# !tensorboard dev upload --logdir runs

# + id="qS5Q9hT2Nwx6" colab={"base_uri": "https://localhost:8080/", "height": 1000} executionInfo={"status": "ok", "timestamp": 1628504218136, "user_tz": -120, "elapsed": 2862968, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="f056b058-a0ce-44a9-e649-263ca8d64381"
# %%time
trainer.train()

# + id="MdHQ-R8XNwx6" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628504219031, "user_tz": -120, "elapsed": 911, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="a1cb9adf-0c4b-4e9e-b50a-b65fb9bdbbdf"
trainer.save_model(model_path)

# + id="MeNvtqwLfj4H" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628504219032, "user_tz": -120, "elapsed": 24, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="e65f364a-6cd3-4450-e35f-957b17cbb364"
# !ls -l --block-size=M {model_path}

# + colab={"base_uri": "https://localhost:8080/"} id="NivCzVmChYft" executionInfo={"status": "ok", "timestamp": 1628504219033, "user_tz": -120, "elapsed": 13, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="3bc097d9-5e9d-4f6c-80bb-335c1c00cec0"
# !mkdir drive/MyDrive/models/{model_path}

# + id="ueXkcK7rf5SI" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1628504219527, "user_tz": -120, "elapsed": 500, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gj7ldLv317JVrv6jtloGcYIgoqAhT08NCy6q5wtrA=s64", "userId": "06955205838397477388"}} outputId="dbbbbcc5-0029-4569-a535-7c19d5b6faf4"
# !cp {model_path}/* drive/MyDrive/models/{model_path}/

StarcoderdataJupyterScriptsDedupFiltered