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

# + id="o_0K1lsW1dj9"
"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.

Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""
# If you're using Google Colab and not running locally, run this cell

# install NeMo
BRANCH = 'v1.0.0b2'
# !python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]



# + pycharm={"name": "#%%\n"} id="95FHWXOVpUFp"
# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:
# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'
# ! pip install ipywidgets
# ! jupyter nbextension enable --py widgetsnbextension

# Please restart the kernel after running this cell

# + id="dzqD2WDFOIN-"
from nemo.collections import nlp as nemo_nlp
from nemo.utils.exp_manager import exp_manager

import os
import wget 
import torch
import pytorch_lightning as pl
from omegaconf import OmegaConf

# + [markdown] id="daYw_Xll2ZR9"
# # Task Description
# **Sentiment Analysis** is the task of detecting the sentiment in text. We model this problem as a simple form of a text classification problem. For example `Gollum's performance is incredible!` has a positive sentiment while `It's neither as romantic nor as thrilling as it should be.` has a negative sentiment.
# .

# + [markdown] id="ZnuziSwJ1yEB"
# # Dataset
#
# In this tutorial we going to use [The Stanford Sentiment Treebank (SST-2)](https://nlp.stanford.edu/sentiment/index.html) corpus for sentiment analysis. This version of the dataset contains a collection of sentences with binary labels of positive and negative. It is a standard benchmark for sentence classification and is part of the GLUE Benchmark: https://gluebenchmark.com/tasks. Please download and unzip the SST-2 dataset from GLUE. It should contain three files of train.tsv, dev.tsv, and test.tsv which can be used for training, validation, and test respectively.
#
#
#
#

# + [markdown] id="qzcZ3nb_-SVT"
# # NeMo Text Classification Data Format
#
# [TextClassificationModel](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/text_classification/text_classification_model.py) in NeMo supports text classification problems such as sentiment analysis or domain/intent detection for dialogue systems, as long as the data follows the format specified below. 
#
# TextClassificationModel requires the data to be stored in TAB separated files (.tsv) with two columns of sentence and label. Each line of the data file contains text sequences, where words are separated with spaces and label separated with [TAB], i.e.: 
#
# ```
# [WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL]
# ```
#
# For example:
# ```
# hide new secretions from the parental units[TAB]0
#
# that loves its characters and communicates something rather beautiful about human nature[TAB]1
# ...
# ```
#
#
# If your dataset is stored in another format, you need to convert it to this format to use the TextClassificationModel. 

# + [markdown] id="SL58EWkd2ZVb"
# ## Download and Preprocess the Data
#
# First, you need to download the zipped file of the SST-2 dataset from the GLUE Benchmark website: https://gluebenchmark.com/tasks, and put it in the current folder. Then the following script would extract it into the data path specified by `DATA_DIR`:

# + id="n8HZrDmr12_-"
DATA_DIR = "DATA_DIR"
WORK_DIR = "WORK_DIR"
os.environ['DATA_DIR'] = DATA_DIR

os.makedirs(WORK_DIR, exist_ok=True)
os.makedirs(DATA_DIR, exist_ok=True)

# ! unzip -o SST-2.zip -d {DATA_DIR}


# + [markdown] id="U8Ty5_S7Ye8h"
# Now, the data folder should contain the following files:

# + [markdown] id="L8vsyh3JZH26"
#
#
# * train.tsv
# * dev.tsv
# * test.tsv
#
#
# The format of `train.tsv` and `dev.tsv` is close to NeMo's format except to have an extra header line at the beginning of the files. We would remove these extra lines. But `test.tsv` has different format and labels are missing for this part of the data.
#

# + id="qB0oLE4R9EhJ"
# ! sed 1d {DATA_DIR}/SST-2/train.tsv > {DATA_DIR}/SST-2/train_nemo_format.tsv
# ! sed 1d {DATA_DIR}/SST-2/dev.tsv > {DATA_DIR}/SST-2/dev_nemo_format.tsv
# ! ls -l {DATA_DIR}/SST-2

# + id="6UDPgadLN6SG"
# let's take a look at the data 
print('Contents (first 5 lines) of train.tsv:')
# ! head -n 5 {DATA_DIR}/SST-2/train_nemo_format.tsv

print('\nContents (first 5 lines) of test.tsv:')
# ! head -n 5 {DATA_DIR}/SST-2/test.tsv

# + [markdown] id="daludzzL2Jba"
# # Model Configuration

# + [markdown] id="_whKCxfTMo6Y"
# Now, let's take a closer look at the model's configuration and learn to train the model from scratch and finetune the pretrained model.
#
# Our text classification model uses a pretrained [BERT](https://arxiv.org/pdf/1810.04805.pdf) model (or other BERT-like models) followed by a classification layer on the output of the first token ([CLS]).
#
# The model is defined in a config file which declares multiple important sections. The most important ones are:
# - **model**: All arguments that are related to the Model - language model, tokenizer, head classifier, optimizer, schedulers, and datasets/data loaders.
#
# - **trainer**: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, precision level, etc.
#

# + id="T1gA8PsJ13MJ"
# download the model's configuration file 
MODEL_CONFIG = "text_classification_config.yaml"
CONFIG_DIR = WORK_DIR + '/configs/'

os.makedirs(CONFIG_DIR, exist_ok=True)
if not os.path.exists(CONFIG_DIR + MODEL_CONFIG):
    print('Downloading config file...')
    wget.download('https://raw.githubusercontent.com/NVIDIA/NeMo/v1.0.0b2/examples/nlp/text_classification/conf/' + MODEL_CONFIG, CONFIG_DIR)
    print('Config file downloaded!')
else:
    print ('config file already exists')
config_path = f'{WORK_DIR}/configs/{MODEL_CONFIG}'
print(config_path)
config = OmegaConf.load(config_path)

# + [markdown] id="mX3KmWMvSUQw"
# ## this line will print the entire config of the model
# print(OmegaConf.to_yaml(config))

# + [markdown] id="ZCgWzNBkaQLZ"
# # Model Training From Scratch
# ## Setting up data within the config
#
# We first need to set the num_classes in the config file which specifies the number of classes in the dataset. For SST-2, we have just two classes (0-positive and 1-negative). So we set the num_classes to 2. The model supports more than 2 classes too.
#
#

# + id="jFSMiWtlkaC5"
config.model.dataset.num_classes=2

# + [markdown] id="xCkc7QGikqPh"
#
# Among other things, the config file contains dictionaries called dataset, train_ds and validation_ds. These are configurations used to setup the Dataset and DataLoaders of the corresponding config.
#
# Notice that some config lines, including `model.dataset.classes_num`, have `???` as their value, this means that values for these fields are required to be to be specified by the user. We need to specify and set the `model.train_ds.file_name`, `model.validation_ds.file_name`, and `model.test_ds.file_name` in the config file to the paths of the train, validation, and test files if they exist. We may do it by updating the config file or by setting them from the command line. 
#
# Let's now set the train and validation paths in the config.

# + id="LQHCJN-ZaoLp"
config.model.train_ds.file_path = os.path.join(DATA_DIR, 'SST-2/train_nemo_format.tsv')
config.model.validation_ds.file_path = os.path.join(DATA_DIR, 'SST-2/dev_nemo_format.tsv')
# Name of the .nemo file where trained model will be saved.
config.save_to = 'trained-model.nemo'
config.export_to = 'trained-model.onnx'

print("Train dataloader's config: \n")
# OmegaConf.to_yaml() is used to create a proper format for printing the train dataloader's config
# You may change other params like batch size or the number of samples to be considered (-1 means all the samples)
print(OmegaConf.to_yaml(config.model.train_ds))

# + [markdown] id="<KEY>"
# ## Building the PyTorch Lightning Trainer
#
# NeMo models are primarily PyTorch Lightning (PT) modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem.
#
# Let's first instantiate a PT Trainer object by using the trainer section of the config.

# + id="1tG4FzZ4Ui60"
print("Trainer config - \n")
# OmegaConf.to_yaml() is used to create a proper format for printing the trainer config
print(OmegaConf.to_yaml(config.trainer))

# + [markdown] id="XMVVs0INi5zj"
# First you need to create a PT trainer with the params stored in the trainer's config. You may set the number of steps for training with max_steps or number of epochs with max_epochs in the trainer's config.

# + id="knF6QeQQdMrH"
# lets modify some trainer configs
# checks if we have GPU available and uses it
config.trainer.gpus = 1 if torch.cuda.is_available() else 0


# for mixed precision training, uncomment the lines below (precision should be set to 16 and amp_level to O1):
# config.trainer.precision = 16
# config.trainer.amp_level = O1

# disable distributed training when using Colab to prevent the errors
config.trainer.accelerator = None

# setup max number of steps to reduce training time for demonstration purposes of this tutorial
# Training stops when max_step or max_epochs is reached (earliest)
config.trainer.max_epochs = 5

# instantiates a PT Trainer object by using the trainer section of the config
trainer = pl.Trainer(**config.trainer)

# + [markdown] id="8IlEMdVxdr6p"
# ## Setting up the NeMo Experiment¶
#

# + [markdown] id="6Kl5IdnV3O8y"
# NeMo has an experiment manager that handles the logging and saving checkpoints for us, so let's setup it. We need the PT trainer and the exp_manager config:

# + id="8uztqGAmdrYt"
# The experiment manager of a trainer object can not be set twice. We repeat the trainer creation code again here to prevent getting error when this cell is executed more than once. 
trainer = pl.Trainer(**config.trainer)

# exp_dir specifies the path to store the the checkpoints and also the logs, it's default is "./nemo_experiments"
# You may set it by uncommentig the following line
# config.exp_manager.exp_dir = 'LOG_CHECKPOINT_DIR'

# OmegaConf.to_yaml() is used to create a proper format for printing the trainer config
print(OmegaConf.to_yaml(config.exp_manager))

exp_dir = exp_manager(trainer, config.exp_manager)

# the exp_dir provides a path to the current experiment for easy access
print(exp_dir)


# + [markdown] id="8tjLhUvL_o7_"
# Before initializing the model, we might want to modify some of the model configs. For example, we might want to modify the pretrained BERT model to another model. The default model is `bert-base-uncased`. We support a variety of models including all the models available in HuggingFace, and Megatron.

# + id="Xeuc2i7Y_nP5"
# complete list of supported BERT-like models
print(nemo_nlp.modules.get_pretrained_lm_models_list())



# + id="RK2xglXyAUOO"
# specify the BERT-like model, you want to use
# set the `model.language_modelpretrained_model_name' parameter in the config to the model you want to use
config.model.language_model.pretrained_model_name = "bert-base-uncased"

# + [markdown] id="fzNZNAVRjDD-"
# Now, we are ready to initialize our model. During the model initialization call, the dataset and data loaders will also be prepared for the training and validation.
#
# Also, the pretrained BERT model will be automatically downloaded. Note it can take up to a few minutes depending on the size of the chosen BERT model for the first time you create the model. If your dataset is large, it also may take some time to read and process all the datasets. 
#
# Now we can create the model with the model config and the trainer object like this:

# + id="NgsGLydWo-6-"
model = nemo_nlp.models.TextClassificationModel(cfg=config.model, trainer=trainer)

# + [markdown] id="kQ592Tx4pzyB"
# ## Monitoring Training Progress
# Optionally, you can create a Tensorboard visualization to monitor training progress.

# + id="mTJr16_pp0aS"
try:
  from google import colab
  COLAB_ENV = True
except (ImportError, ModuleNotFoundError):
  COLAB_ENV = False

# Load the TensorBoard notebook extension
if COLAB_ENV:
  # %load_ext tensorboard
  # %tensorboard --logdir {exp_dir}
else:
  print("To use tensorboard, please use this notebook in a Google Colab environment.")


# + [markdown] id="MW_JVIi5z68e"
# ## Training
#
# You may start the training by using the trainer.fit() method. The number of steps/epochs of the training are specified already in the config of the trainer and you may update them before creating the trainer.

# + id="hUvnSpyjp0Dh"
# start model training
trainer.fit(model)
model.save_to(config.save_to)

# + [markdown] id="VPdzJVAgSFaJ"
# # Evaluation
#
# To see how the model performs, we can run evaluate and test the performance of the trained model on a data file. Here we would load the best checkpoint (the one with the lowest validation loss) and create a model (eval_model) from the checkpoint. We would also create a new trainer (eval_trainer) to show how it is done when training is done and you have just the checkpoints. If you want to perform the evaluation in the same script as the training's script, you may still use the same model and trainer you used for training.

# + id="92PB0iTqNnW-"
# extract the path of the best checkpoint from the training, you may update it to any checkpoint
checkpoint_path = trainer.checkpoint_callback.best_model_path
# Create an evaluation model and load the checkpoint
eval_model = nemo_nlp.models.TextClassificationModel.load_from_checkpoint(checkpoint_path=checkpoint_path)

# create a dataloader config for evaluation, the same data file provided in validation_ds is used here
# file_path can get updated with any file
eval_config = OmegaConf.create({'file_path': config.model.validation_ds.file_path, 'batch_size': 64, 'shuffle': False, 'num_samples': -1})
eval_model.setup_test_data(test_data_config=eval_config)
#eval_dataloader = eval_model._create_dataloader_from_config(cfg=eval_config, mode='test')

# a new trainer is created to show how to evaluate a checkpoint from an already trained model
# create a copy of the trainer config and update it to be used for final evaluation
eval_trainer_cfg = config.trainer.copy()
eval_trainer_cfg.gpus = 1 if torch.cuda.is_available() else 0 # it is safer to perform evaluation on single GPU as PT is buggy with the last batch on multi-GPUs
eval_trainer_cfg.accelerator = None # 'ddp' is buggy with test process in the current PT, it looks like it has been fixed in the latest master
eval_trainer = pl.Trainer(**eval_trainer_cfg)

eval_trainer.test(model=eval_model, verbose=False) # test_dataloaders=eval_dataloader,


# + [markdown] id="Tit5kG4Z5SXu"
# # Inference
#
# You may create a model from a saved chechpoint and use the model.infer() method to perform inference on a list of queries. There is no need of any trainer for inference.

# + id="BKe5Jn4u9xng"
# extract the path of the best checkpoint from the training, you may update it to any other checkpoint file
checkpoint_path = trainer.checkpoint_callback.best_model_path
# Create an evaluation model and load the checkpoint
infer_model = nemo_nlp.models.TextClassificationModel.load_from_checkpoint(checkpoint_path=checkpoint_path)


# + [markdown] id="y8SFxPJd-hkH"
# To see how the model performs, let’s get model's predictions for a few examples:

# + id="DQhsamclRtxJ"
# move the model to the desired device for inference
# we move the model to "cuda" if available otherwise "cpu" would be used
if torch.cuda.is_available():
    infer_model.to("cuda")
else:
    infer_model.to("cpu")

# define the list of queries for inference
queries = ['by the end of no such thing the audience , like beatrice , has a watchful affection for the monster .', 
           'director <NAME> went out gunning to make a great one .', 
           'uneasy mishmash of styles and genres .']
           
# max_seq_length=512 is the maximum length BERT supports.       
results = infer_model.classifytext(queries=queries, batch_size=3, max_seq_length=512)

print('The prediction results of some sample queries with the trained model:')
for query, result in zip(queries, results):
    print(f'Query : {query}')
    print(f'Predicted label: {result}')


# + [markdown] id="ref1qSonGNhP"
# ## Training Script
#
# If you have NeMo installed locally (eg. cloned from the Github), you can also train the model with `examples/nlp/text_classification/text_classifciation_with_bert.py`. This script contains an example on how to train, evaluate and perform inference with the TextClassificationModel.
#
# For example the following would train a model for 50 epochs in 2 GPUs on a classification task with 2 classes:
#
# ```
# # python text_classification_with_bert.py 
#         model.dataset.num_classes=2
#         model.train_ds=PATH_TO_TRAIN_FILE 
#         model.validation_ds=PATH_TO_VAL_FILE  
#         trainer.max_epochs=50
#         trainer.gpus=2
# ```
#
# This script would also reload the best checkpoint after the training is done and does evaluation on the dev set. Then perform inference on some sample queries. 
#
#
# By default, this script uses `examples/nlp/text_classification/conf/text_classifciation_config.py` config file, and you may update all the params in the config file from the command line. You may also use another config file like this:
#
# ```
# # python text_classification_with_bert.py --config-name==PATH_TO_CONFIG_FILE
#         model.dataset.num_classes=2
#         model.train_ds=PATH_TO_TRAIN_FILE 
#         model.validation_ds=PATH_TO_VAL_FILE  
#         trainer.max_epochs=50
#         trainer.gpus=2
# ```
#
# -

# ## Deployment
#
# You can also deploy a model to an inference engine (like TensorRT or ONNXRuntime) using ONNX exporter.
# If you don't have one, let's install it:

# mkdir ort
# cd ort
git clone --depth 1 --branch v1.5.1 https://github.com/microsoft/onnxruntime.git .
./build.sh --skip_tests --config Release --build_shared_lib --parallel --use_cuda --cuda_home /usr/local/cuda --cudnn_home /usr/lib/x86_64-linux-gnu --build_wheel
pip install ./build/Linux/Release/dist/onnxruntime*.whl
# cd ..

# Then export

model.export(config.export_to)

# And run some queries

# +
import numpy as np
import torch
from nemo.utils import logging
from nemo.collections.nlp.parts.utils_funcs import tensor2list
from nemo.collections.nlp.models.text_classification import TextClassificationModel
from nemo.collections.nlp.data.text_classification import TextClassificationDataset

import onnxruntime

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def postprocessing(results, labels):
    return [labels[str(r)] for r in results]

def create_infer_dataloader(model, queries):
    batch_size = len(queries)
    dataset = TextClassificationDataset(tokenizer=model.tokenizer, queries=queries, max_seq_length=512)
    return torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=2,
        pin_memory=True,
        drop_last=False,
        collate_fn=dataset.collate_fn,
    )


queries = ["by the end of no such thing the audience , like beatrice , has a watchful affection for the monster .",
"director rob marshall went out gunning to make a great one .",
"uneasy mishmash of styles and genres .",
"I love exotic science fiction / fantasy movies but this one was very unpleasant to watch . Suggestions and images of child abuse , mutilated bodies (live or dead) , other gruesome scenes , plot holes , boring acting made this a regretable experience , The basic idea of entering another person's mind is not even new to the movies or TV (An Outer Limits episode was better at exploring this idea) . i gave it 4 / 10 since some special effects were nice ."]

model.eval()
infer_datalayer = create_infer_dataloader(model, queries)

ort_session = onnxruntime.InferenceSession(config.export_to)

for batch in infer_datalayer:
    input_ids, input_type_ids, input_mask, subtokens_mask = batch
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_ids),
                  ort_session.get_inputs()[1].name: to_numpy(input_mask),
                  ort_session.get_inputs()[2].name: to_numpy(input_type_ids),}
    ologits = ort_session.run(None, ort_inputs)
    alogits = np.asarray(ologits)
    logits = torch.from_numpy(alogits[0])
    preds = tensor2list(torch.argmax(logits, dim = -1))
    processed_results = postprocessing(preds, {"0": "negative", "1": "positive"})

    logging.info('The prediction results of some sample queries with the trained model:')
    for query, result in zip(queries, processed_results):
        logging.info(f'Query : {query}')
        logging.info(f'Predicted label: {result}')
    break

StarcoderdataJupyterScriptsDedupFiltered