#!/usr/bin/env python
# coding: utf-8

# # Programming Language Classification with OpenVINO
# ## Overview
# This tutorial will be divided in 2 parts:
# 1. Create a simple inference pipeline with a pre-trained model using the OpenVINO™ IR format.
# 2. Conduct [post-training quantization](https://docs.openvino.ai/2023.3/ptq_introduction.html) on a pre-trained model using Hugging Face Optimum and benchmark performance.
# 
# Feel free to use the notebook outline in Jupyter or your IDE for easy navigation. 
# 
# #### Table of contents:
# - [Introduction](#Introduction)
#     - [Task](#Task)
#     - [Model](#Model)
# - [Part 1: Inference pipeline with OpenVINO](#Part-1:-Inference-pipeline-with-OpenVINO)
#     - [Install prerequisites](#Install-prerequisites)
#     - [Imports](#Imports)
#     - [Setting up HuggingFace cache](#Setting-up-HuggingFace-cache)
#     - [Select inference device](#Select-inference-device)
#     - [Download resources](#Download-resources)
#     - [Create inference pipeline](#Create-inference-pipeline)
#     - [Inference on new input](#Inference-on-new-input)
# - [Part 2: OpenVINO post-training quantization with HuggingFace Optimum](#Part-2:-OpenVINO-post-training-quantization-with-HuggingFace-Optimum)
#     - [Define constants and functions](#Define-constants-and-functions)
#     - [Load resources](#Load-resources)
#     - [Load calibration dataset](#Load-calibration-dataset)
#     - [Quantize model](#Quantize-model)
#     - [Load quantized model](#Load-quantized-model)
#     - [Inference on new input using quantized model](#Inference-on-new-input-using-quantized-model)
#     - [Load evaluation set](#Load-evaluation-set)
#     - [Evaluate model](#Evaluate-model)
# - [Additional resources](#Additional-resources)
# - [Clean up](#Clean-up)
# 

# ## Introduction
# [back to top ⬆️](#Table-of-contents:)
# 
# ### Task
# [back to top ⬆️](#Table-of-contents:)
# 
# **Programming language classification** is the task of identifying which programming language is used in an arbitrary code snippet. This can be useful to label new data to include in a dataset, and potentially serve as an intermediary step when input snippets need to be process based on their programming language.
# 
# It is a relatively easy machine learning task given that each programming language has its own formal symbols, syntax, and grammar. However, there are some potential edge cases:
# - **Ambiguous short snippets**: For example, TypeScript is a superset of JavaScript, meaning it does everything JavaScript can and more. For a short input snippet, it might be impossible to distinguish between the two. Given we know TypeScript is a superset, and the model doesn't, we should default to classifying the input as JavaScript in a post-processing step.
# - **Nested programming languages**: Some languages are typically used in tandem. For example, most HTML contains CSS and JavaScript, and it is not uncommon to see SQL nested in other scripting languages. For such input, it is unclear what the expected output class should be. 
# - **Evolving programming language**: Even though programming languages are formal, their symbols, syntax, and grammar can be revised and updated. For example, the walrus operator (`:=`) was a symbol distinctively used in Golang, but was later introduced in Python 3.8.

# ### Model
# [back to top ⬆️](#Table-of-contents:)
# 
# The classification model that will be used in this notebook is [`CodeBERTa-language-id`](https://huggingface.co/huggingface/CodeBERTa-language-id) by HuggingFace. This model was fine-tuned from the masked language modeling model [`CodeBERTa-small-v1`](https://huggingface.co/huggingface/CodeBERTa-small-v1) trained on the [`CodeSearchNet`](https://huggingface.co/huggingface/CodeBERTa-small-v1) dataset (Husain, 2019).
# 
# It supports 6 programming languages:
# - Go
# - Java
# - JavaScript
# - PHP
# - Python
# - Ruby

# ## Part 1: Inference pipeline with OpenVINO
# [back to top ⬆️](#Table-of-contents:)
# 
# For this section, we will use the [**HuggingFace Optimum**](https://huggingface.co/docs/optimum/index) library, which aims to optimize inference on specific hardware and integrates with the OpenVINO toolkit. The code will be very similar to the [**HuggingFace Transformers**](https://huggingface.co/docs/transformers/index), but will allow to automatically convert models to the OpenVINO™ IR format.

# ### Install prerequisites
# [back to top ⬆️](#Table-of-contents:)
# 
# First, complete the [repository installation steps](../../README.md).
# 
# Then, the following cell will install:
# - HuggingFace Optimum with OpenVINO support
# - HuggingFace Evaluate to benchmark results

# In[1]:


get_ipython().run_line_magic('pip', 'install -q "diffusers>=0.17.1" "openvino>=2023.1.0" "nncf>=2.5.0" "gradio" "onnx>=1.11.0" "transformers>=4.33.0" "evaluate" --extra-index-url https://download.pytorch.org/whl/cpu')
get_ipython().run_line_magic('pip', 'install -q "git+https://github.com/huggingface/optimum-intel.git"')


# ### Imports
# [back to top ⬆️](#Table-of-contents:)
# 
# The import `OVModelForSequenceClassification` from Optimum is equivalent to `AutoModelForSequenceClassification` from Transformers

# In[2]:


from functools import partial
from pathlib import Path

import pandas as pd
from datasets import load_dataset, Dataset
import evaluate
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from optimum.intel import OVModelForSequenceClassification  
from optimum.intel.openvino import OVConfig, OVQuantizer
from huggingface_hub.utils import RepositoryNotFoundError


# ### Setting up HuggingFace cache
# [back to top ⬆️](#Table-of-contents:)
# 
# Resources from HuggingFace will be downloaded in the local folder `./model` (next to this notebook) instead of the device global cache for easy cleanup. Learn more [here](https://huggingface.co/docs/transformers/installation?highlight=transformers_cache#cache-setup).

# In[3]:


MODEL_NAME = "CodeBERTa-language-id"
MODEL_ID = f"huggingface/{MODEL_NAME}"
MODEL_LOCAL_PATH = Path("./model").joinpath(MODEL_NAME)


# ### Select inference device
# [back to top ⬆️](#Table-of-contents:)
# 
# select device from dropdown list for running inference using OpenVINO

# In[4]:


import ipywidgets as widgets
import openvino as ov

core = ov.Core()

device = widgets.Dropdown(
    options=core.available_devices + ["AUTO"],
    value='AUTO',
    description='Device:',
    disabled=False,
)

device


# ### Download resources
# [back to top ⬆️](#Table-of-contents:)
# 

# In[5]:


# try to load resources locally
try:
    model = OVModelForSequenceClassification.from_pretrained(MODEL_LOCAL_PATH, device=device.value)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_LOCAL_PATH)
    print(f"Loaded resources from local path: {MODEL_LOCAL_PATH.absolute()}")

# if not found, download from HuggingFace Hub then save locally
except (RepositoryNotFoundError, OSError):
    print("Downloading resources from HuggingFace Hub")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    tokenizer.save_pretrained(MODEL_LOCAL_PATH)

    # export=True is needed to convert the PyTorch model to OpenVINO
    model = OVModelForSequenceClassification.from_pretrained(MODEL_ID, export=True, device=device.value)
    model.save_pretrained(MODEL_LOCAL_PATH)
    print(f"Ressources cached locally at: {MODEL_LOCAL_PATH.absolute()}")


# ### Create inference pipeline
# [back to top ⬆️](#Table-of-contents:)
# 

# In[6]:


code_classification_pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)


# ### Inference on new input
# [back to top ⬆️](#Table-of-contents:)
# 

# In[7]:


# change input snippet to test model
input_snippet = "df['speed'] = df.distance / df.time"
output = code_classification_pipe(input_snippet)

print(f"Input snippet:\n  {input_snippet}\n")
print(f"Predicted label: {output[0]['label']}")
print(f"Predicted score: {output[0]['score']:.2}")


# ## Part 2: OpenVINO post-training quantization with HuggingFace Optimum
# [back to top ⬆️](#Table-of-contents:)
# 
# In this section, we will quantize a trained model. At a high-level, this process consists of using lower precision numbers in the model, which results in a smaller model size and faster inference at the cost of a potential marginal performance degradation. [Learn more](https://docs.openvino.ai/2023.3/ptq_introduction.html).
# 
# The HuggingFace Optimum library supports post-training quantization for OpenVINO. [Learn more](https://huggingface.co/docs/optimum/main/en/intel/index).

# ### Define constants and functions
# [back to top ⬆️](#Table-of-contents:)
# 

# In[8]:


QUANTIZED_MODEL_LOCAL_PATH = MODEL_LOCAL_PATH.with_name(f"{MODEL_NAME}-quantized")
DATASET_NAME = "code_search_net"
LABEL_MAPPING = {"go": 0, "java": 1, "javascript": 2, "php": 3, "python": 4, "ruby": 5}


def preprocess_function(examples: dict, tokenizer):
    """Preprocess inputs by tokenizing the `func_code_string` column"""
    return tokenizer(
        examples["func_code_string"],
        padding="max_length",
        max_length=tokenizer.model_max_length,
        truncation=True,
    )


def map_labels(example: dict) -> dict:
    """Convert string labels to integers"""
    label_mapping = {"go": 0, "java": 1, "javascript": 2, "php": 3, "python": 4, "ruby": 5}
    example["language"] = label_mapping[example["language"]]
    return example 


def get_dataset_sample(dataset_split: str, num_samples: int) -> Dataset:
    """Create a sample with equal representation of each class without downloading the entire data"""
    labels = ["go", "java", "javascript", "php", "python", "ruby"]
    example_per_label = num_samples // len(labels)

    examples = []
    for label in labels:
        subset = load_dataset("code_search_net", split=dataset_split, name=label, streaming=True)
        subset = subset.map(map_labels)
        examples.extend([example for example in subset.shuffle().take(example_per_label)])
    
    return Dataset.from_list(examples)


# ### Load resources
# [back to top ⬆️](#Table-of-contents:)
# 
# NOTE: the base model is loaded using `AutoModelForSequenceClassification` from `Transformers`

# In[9]:


tokenizer = AutoTokenizer.from_pretrained(MODEL_LOCAL_PATH)
base_model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

quantizer = OVQuantizer.from_pretrained(base_model)
quantization_config = OVConfig()


# ### Load calibration dataset
# [back to top ⬆️](#Table-of-contents:)
# 
# The `get_dataset_sample()` function will sample up to `num_samples`, with an equal number of examples across the 6 programming languages.
# 
# NOTE: Uncomment the method below to download and use the full dataset (5+ Gb). 

# In[10]:


calibration_sample = get_dataset_sample(dataset_split="train", num_samples=120)
calibration_sample = calibration_sample.map(partial(preprocess_function, tokenizer=tokenizer))

# calibration_sample = quantizer.get_calibration_dataset(
#     DATASET_NAME,
#     preprocess_function=partial(preprocess_function, tokenizer=tokenizer),
#     num_samples=120,
#     dataset_split="train",
#     preprocess_batch=True,
# )


# ### Quantize model
# [back to top ⬆️](#Table-of-contents:)
# 
# Calling `quantizer.quantize(...)` will iterate through the calibration dataset to quantize and save the model

# In[11]:


quantizer.quantize(
    quantization_config=quantization_config,
    calibration_dataset=calibration_sample,
    save_directory=QUANTIZED_MODEL_LOCAL_PATH,
)


# ### Load quantized model
# [back to top ⬆️](#Table-of-contents:)
# 
# NOTE: the argument `export=True` is not required since the quantized model is already in the OpenVINO format.

# In[12]:


quantized_model = OVModelForSequenceClassification.from_pretrained(QUANTIZED_MODEL_LOCAL_PATH, device=device.value)
quantized_code_classification_pipe = pipeline("text-classification", model=quantized_model, tokenizer=tokenizer)


# ### Inference on new input using quantized model
# [back to top ⬆️](#Table-of-contents:)
# 

# In[13]:


input_snippet = "df['speed'] = df.distance / df.time"
output = quantized_code_classification_pipe(input_snippet)

print(f"Input snippet:\n  {input_snippet}\n")
print(f"Predicted label: {output[0]['label']}")
print(f"Predicted score: {output[0]['score']:.2}")


# ### Load evaluation set
# [back to top ⬆️](#Table-of-contents:)
# 
# NOTE: Uncomment the method below to download and use the full dataset (5+ Gb). 

# In[14]:


validation_sample = get_dataset_sample(dataset_split="validation", num_samples=120)

# validation_sample = load_dataset(DATASET_NAME, split="validation")


# ### Evaluate model
# [back to top ⬆️](#Table-of-contents:)
# 

# In[15]:


# This class is needed due to a current limitation of the Evaluate library with multiclass metrics
# ref: https://discuss.huggingface.co/t/combining-metrics-for-multiclass-predictions-evaluations/21792/16
class ConfiguredMetric:
    def __init__(self, metric, *metric_args, **metric_kwargs):
        self.metric = metric
        self.metric_args = metric_args
        self.metric_kwargs = metric_kwargs
    
    def add(self, *args, **kwargs):
        return self.metric.add(*args, **kwargs)
    
    def add_batch(self, *args, **kwargs):
        return self.metric.add_batch(*args, **kwargs)

    def compute(self, *args, **kwargs):
        return self.metric.compute(*args, *self.metric_args, **kwargs, **self.metric_kwargs)

    @property
    def name(self):
        return self.metric.name

    def _feature_names(self):
        return self.metric._feature_names()


# First, an `Evaluator` object for `text-classification` and a set of `EvaluationModule` are instantiated. Then, the evaluator `.compute()` method is called on both the base `code_classification_pipe` and the quantized `quantized_code_classification_pipeline`. Finally, results are displayed.

# In[16]:


code_classification_evaluator = evaluate.evaluator("text-classification")
# instantiate an object that can contain multiple `evaluate` metrics
metrics = evaluate.combine([
    ConfiguredMetric(evaluate.load('f1'), average='macro'),
])

base_results = code_classification_evaluator.compute(
    model_or_pipeline=code_classification_pipe,
    data=validation_sample,
    input_column="func_code_string",
    label_column="language",
    label_mapping=LABEL_MAPPING,
    metric=metrics,
)

quantized_results = code_classification_evaluator.compute(
    model_or_pipeline=quantized_code_classification_pipe,
    data=validation_sample,
    input_column="func_code_string",
    label_column="language",
    label_mapping=LABEL_MAPPING,
    metric=metrics,
)

results_df = pd.DataFrame.from_records([base_results, quantized_results], index=["base", "quantized"])
results_df


# ## Additional resources
# [back to top ⬆️](#Table-of-contents:)
# - [Grammatical Error Correction with OpenVINO](https://github.com/openvinotoolkit/openvino_notebooks/blob/main/notebooks/214-grammar-correction/214-grammar-correction.ipynb)
# - [Quantize a Hugging Face Question-Answering Model with OpenVINO](https://github.com/huggingface/optimum-intel/blob/main/notebooks/openvino/question_answering_quantization.ipynb)**

# ## Clean up
# [back to top ⬆️](#Table-of-contents:)
# 
# Uncomment and run cell below to delete all resources cached locally in ./model

# In[17]:


# import os
# import shutil

# try:
#     shutil.rmtree(path=QUANTIZED_MODEL_LOCAL_PATH)
#     shutil.rmtree(path=MODEL_LOCAL_PATH)
#     os.remove(path="./compressed_graph.dot")
#     os.remove(path="./original_graph.dot")
# except FileNotFoundError:
#     print("Directory was already deleted")