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

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/IgnatiusEzeani/IGBONLP/blob/master/training_igbo_bert.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="b8y9VYzYvrrh"
# # **Training Igbo BERT Language Model from Scratch**

# + [markdown] id="EtHOxhbOv-m9"
# 1. Find a dataset

# + colab={"base_uri": "https://localhost:8080/"} id="gcu0G7osCloL" outputId="806e7ddf-63cc-49bb-d14b-8be394bbe346"
# !wget -c https://github.com/IgnatiusEzeani/IGBONLP/raw/master/ig_monoling/text.zip

# + id="WvFWfdb1imv_"
import zipfile
import os


def unzip(zipfilename):
  try:
    with zipfile.ZipFile(zipfilename, 'r') as zip_ref:
      zip_ref.extractall(zipfilename[:-4])
      return f"'{zipfilename}' unzipped!"
  except FileNotFoundError:
    print(f"Cannot find '{zipfilename}' file")

unzip("text.zip")
# !rm text.zip

# + id="Q8Lf6yjYjPOj"
# import os
import shutil
dir_name = "/content/text/"
text=""
for fname in os.listdir(dir_name):
  fname = os.path.join(dir_name, fname)
  with open(fname, "r", encoding="utf8") as datafile:
    text = text+"\n"+datafile.read()

with open("data.txt", "w", encoding="utf8") as datafile:
  datafile.write(text)

shutil.rmtree("text")

# + [markdown] id="d1lT9xMAwPxM"
# 2. Train a tokenizer
#
# We choose to train a byte-level Byte-pair encoding tokenizer (the same as GPT-2), with the same special tokens as RoBERTa. Let’s arbitrarily pick its size to be 52,000.
#
# We recommend training a byte-level BPE (rather than let’s say, a WordPiece tokenizer like BERT) because it will start building its vocabulary from an alphabet of single bytes, so all words will be decomposable into tokens (no more <unk> tokens!).

# + colab={"base_uri": "https://localhost:8080/"} id="lJ6Zz20gvmCx" outputId="469eb7cb-c4a0-4417-8890-f589d502f711"
# We won't need TensorFlow here
# !pip uninstall -y tensorflow
# Install `transformers` from master
# !pip install git+https://github.com/huggingface/transformers
# !pip list | grep -E 'transformers|tokenizers'
# transformers version at notebook update --- 2.11.0
# tokenizers version at notebook update --- 0.8.0rc1

# + id="n-eLh9QCzFiH" colab={"base_uri": "https://localhost:8080/"} outputId="b00f56d9-50a1-4329-8595-cecf7d02e93a"
# %%time 
from pathlib import Path
from tokenizers import ByteLevelBPETokenizer

paths = [str(x) for x in Path(".").glob("**/*.txt")]

# Initialize a tokenizer
tokenizer = ByteLevelBPETokenizer()

# Customize training
tokenizer.train(files=paths, vocab_size=52_000, min_frequency=2, special_tokens=[
    "<s>",
    "<pad>",
    "</s>",
    "<unk>",
    "<mask>",
])

# + id="QbIL0wBTb3Ak"
Path(".").glob("**/*.txt") ##remove later...

# + [markdown] id="Lf8aFiBVzE1o"
# Now let's save files to disk

# + colab={"base_uri": "https://localhost:8080/"} id="udz15UV71dwk" outputId="aaff9dfa-29d6-4315-d2fb-8bd3087b1fcf"
# !mkdir igbo_bert
tokenizer.save_model("igbo_bert")

# + [markdown] id="wxOvsxMg5aJ8"
# What is great is that our tokenizer is optimized for Igbo. Compared to a generic tokenizer trained for English, more native words are represented by a single, unsplit token. Diacritics are encoded natively. We also represent sequences in a more efficient manner.
#
# Here’s how you can use it in tokenizers, including handling the RoBERTa special tokens – of course, you’ll also be able to use it directly from transformers.
#

# + id="IC0tDevg6LdD"
from tokenizers.implementations import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing

tokenizer = ByteLevelBPETokenizer(
    "./igbo_bert/vocab.json",
    "./igbo_bert/merges.txt",
)

# + id="wsz-NdVP6k8u"
tokenizer._tokenizer.post_processor = BertProcessing(
    ("</s>", tokenizer.token_to_id("</s>")),
    ("<s>", tokenizer.token_to_id("<s>")),
)
tokenizer.enable_truncation(max_length=512)

# + colab={"base_uri": "https://localhost:8080/"} id="lkSNqBXz7m0x" outputId="e8973dce-f02f-4494-fc19-79c6952a9141"
tokenizer.encode("Simone kọrọ ya akụkọ nke ya ahu.").tokens

# + [markdown] id="qGkjBbMLAMa7"
# 3. Train a language model from scratch
# Update: This section follows along the run_language_modeling.py script, using our new Trainer directly. Feel free to pick the approach you like best.
#
# We’ll train a RoBERTa-like model, which is a BERT-like with a couple of changes (check the documentation for more details).
#
# As the model is BERT-like, we’ll train it on a task of Masked language modeling, i.e. the predict how to fill arbitrary tokens that we randomly mask in the dataset. This is taken care of by the example script.

# + colab={"base_uri": "https://localhost:8080/"} id="I_rcU-fzAWz-" outputId="dded7a93-0b1c-4674-9895-57acf99893d5"
# Check that we have a GPU
# !nvidia-smi

# + colab={"base_uri": "https://localhost:8080/"} id="3CyCk-6rAaaN" outputId="5585c4e5-c300-449c-8660-ed8930f65ced"
# Check that PyTorch sees it
import torch
torch.cuda.is_available()

# + [markdown] id="YrCJXtD_Bfr2"
# We'll define the following config for the model

# + id="RjR_m2g8Bo-g"
from transformers import RobertaConfig

config = RobertaConfig(
    vocab_size=52_000,
    max_position_embeddings=514,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

# + [markdown] id="GE2ry8PbGLPR"
# Now let's re-create our tokenizer in transformers

# + id="I3eivVi0GNJq"
from transformers import RobertaTokenizerFast

tokenizer = RobertaTokenizerFast.from_pretrained("./igbo_bert", max_len=512)

# + [markdown] id="Yc57oI0zGnJl"
# Finally let's initialize our model.
#
# Important:
#
# As we are training from scratch, we only initialize from a config, not from an existing pretrained model or checkpoint.

# + id="pY933eKSGoyu"
from transformers import RobertaForMaskedLM

model = RobertaForMaskedLM(config=config)

# + colab={"base_uri": "https://localhost:8080/"} id="vq-El35_G0D_" outputId="8f5b5f05-2a9b-43d8-f8c4-80717de8144c"
model.num_parameters()
# => 83 million parameters

# + [markdown] id="DUNJL-fTG_bL"
# Now let's build our training Dataset
# We'll build our dataset by applying our tokenizer to our text file.
#
# Here, as we only have one text file, we don't even need to customize our Dataset. We'll just use the LineByLineDataset out-of-the-box.

# + colab={"base_uri": "https://localhost:8080/"} id="Lw72at5ZHKXY" outputId="dfd37de4-9a31-4f8f-a454-97b1b35351c8"
# %%time
from transformers import LineByLineTextDataset

dataset = LineByLineTextDataset(
    tokenizer = tokenizer,
    file_path = "/content/data.txt",
    block_size = 128
)

# + [markdown] id="guy74crmSnoT"
# Like in the run_language_modeling.py script, we need to define a data_collator.
#
# This is just a small helper that will help us batch different samples of the dataset together into an object that PyTorch knows how to perform backprop on.

# + id="PTJiFWlpShfa"
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

# + [markdown] id="sfH1dXhMTxqL"
# Finally, we are all set to initialize our Trainer

# + id="ikFCJiYVTzpo"
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./igbo_bert",
    overwrite_output_dir=True,
    num_train_epochs=2,
    per_gpu_train_batch_size=64,
    save_steps=10_000,
    save_total_limit=2,
    prediction_loss_only=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
)

# + [markdown] id="HpywHdSvURvv"
# Start training

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="Fs6L0X0RUTAH" outputId="c1c84111-2e76-42aa-8759-36b4f1d11edd"
# %%time
trainer.train()

# + [markdown] id="ZlFtQHFlYa6P"
#  Save final model (+ tokenizer + config) to disk

# + colab={"base_uri": "https://localhost:8080/"} id="t_jFqUHCYgXQ" outputId="2add4e51-d4ea-48a9-a098-63727a662c85"
trainer.save_model("./igbo_bert")

# + [markdown] id="RmTbUurZY0dZ"
# # 4. Check that the LM actually trained

# + [markdown] id="UKTL3b5gZCyk"
# Aside from looking at the training and eval losses going down, the easiest way to check whether our language model is learning anything interesting is via the FillMaskPipeline.
#
# Pipelines are simple wrappers around tokenizers and models, and the 'fill-mask' one will let you input a sequence containing a masked token (here, <mask>) and return a list of the most probable filled sequences, with their probabilities.

# + [markdown] id="glwKBga6ZWjz"
#

# + colab={"base_uri": "https://localhost:8080/"} id="xMDIwVOjZjB7" outputId="44fc0d97-fdda-4d2c-c9b5-c85c0d90e6a5"
from transformers import pipeline

fill_mask = pipeline(
    "fill-mask",
    model="./igbo_bert",
    tokenizer="./igbo_bert"
)

# + colab={"base_uri": "https://localhost:8080/"} id="HrvKuawbZ7bw" outputId="47272296-41b1-4967-956a-d84d3885631f"
# The sun <mask>.
# =>

fill_mask("Uwa anyị niile <mask>.") #=eri
# fill_mask("Nwaanyị na <mask> ji na akara.") #=eri

# + [markdown] id="J_FYJ6rcaDMh"
# Ok, simple syntax/grammar works. Let’s try a slightly more interesting prompt:

# + id="JANBv48RaEb7" colab={"base_uri": "https://localhost:8080/"} outputId="b6986153-83d6-4e9c-802c-ed032b299adc"
fill_mask("O riri <mask>.")

# This is the beginning of a beautiful <mask>.
# =>

# + [markdown] id="uK7JNpX2QpG2"
# Zip and download the model file

# + colab={"base_uri": "https://localhost:8080/", "height": 36} id="v0J-GhS65ouL" outputId="e882b147-7bd9-471c-daa5-231a8d8184ed"
shutil.make_archive("/content/igbo_bert", 'zip', "igbo_bert")

# + colab={"base_uri": "https://localhost:8080/", "height": 320} id="AHXdwg9BRAKC" outputId="c4c3d6cf-df88-4df1-bd40-696288eff6df"
from google.colab import files
files.download("/content/igbo_bert.zip")

# + colab={"base_uri": "https://localhost:8080/", "height": 132} id="DJCZUo9UC0Qv" outputId="14a61be6-4d4f-4fb5-8295-57236677b18b"
model_save_name = 'IGBOchy'
path = F"/content/drive/Shareddrives/{model_save_name}" 
torch.save(model.state_dict(), /content/drive/Shareddrives)

# + colab={"base_uri": "https://localhost:8080/"} id="jQi0FXKFG9vO" outputId="27c55d6d-f403-4c34-bae0-218fa5f7a5f6"
# !cp "/content/IGBOchy" "/content/drive/Shareddrives"

# + colab={"base_uri": "https://localhost:8080/"} id="AtFXp_3mIq1w" outputId="13bbcd5e-4dec-4eef-daa3-4a6874ce2cca"
# !ls /content/drive

# + colab={"base_uri": "https://localhost:8080/", "height": 132} id="m6LSEydAOhq5" outputId="ce43ac1a-9443-4b89-eb5c-379e6b7d1b69"
torch.save("IGBOchy", /content/drive/Shareddrives)

# + colab={"base_uri": "https://localhost:8080/"} id="0WvY6mdXK_dv" outputId="5712eda8-4881-42a2-b687-337f1d4f213e"
#reading a text file saved in github using requests package
import requests
url = 'https://github.com/chiamaka249/Igbo-merged-corpus-/blob/main/IGBO_CORPUS_MERGE.zip'
page = requests.get(url)
print (page.text)

StarcoderdataJupyterScriptsDedupFiltered