main.py

"""⭐ Text Classification with Optimum and ONNXRuntime

Streamlit application to classify text using multiple models.

Author:
    - @ChainYo - https://github.com/ChainYo
"""

import plotly
import numpy as np
import pandas as pd
import streamlit as st

from pathlib import Path
from time import sleep
from typing import Dict, List, Union

from optimum.onnxruntime import ORTModelForSequenceClassification, ORTOptimizer, ORTQuantizer
from optimum.onnxruntime.configuration import OptimizationConfig, AutoQuantizationConfig
from optimum.onnxruntime.model import ORTModel
from optimum.pipelines import pipeline as ort_pipeline
from transformers import BertTokenizer, BertForSequenceClassification, pipeline

from utils import calculate_inference_time


HUB_MODEL_PATH = "yiyanghkust/finbert-tone"
BASE_PATH = Path("models")
ONNX_MODEL_PATH = BASE_PATH.joinpath("model.onnx")
OPTIMIZED_BASE_PATH = BASE_PATH.joinpath("optimized")
OPTIMIZED_MODEL_PATH = OPTIMIZED_BASE_PATH.joinpath("model-optimized.onnx")
QUANTIZED_BASE_PATH = BASE_PATH.joinpath("quantized")
QUANTIZED_MODEL_PATH = QUANTIZED_BASE_PATH.joinpath("model-quantized.onnx")
VAR2LABEL = {
    "pt_pipeline": "PyTorch", 
    "ort_pipeline": "ONNXRuntime", 
    "ort_optimized_pipeline": "ONNXRuntime (Optimized)", 
    "ort_quantized_pipeline": "ONNXRuntime (Quantized)",
}


def get_timers(
    samples: Union[List[str], str], exp_number: int, only_mean: bool = False
) -> Dict[str, float]:
    """
    Calculate inference time for each model for a given sample or list of samples.

    Parameters
    ----------
    samples : Union[List[str], str]
        Sample or list of samples to calculate inference time for.
    exp_number : int
        Number of experiments to run.
    
    Returns
    -------
    Dict[str, float]
        Dictionary of inference times for each model for the given samples.
    """
    if isinstance(samples, str):
        samples = [samples]
    
    timers: Dict[str, float] = {}
    for model in VAR2LABEL.keys():
        time_buffer = []
        for _ in range(exp_number):
            with calculate_inference_time(time_buffer):
                st.session_state[model](samples)
        timers[VAR2LABEL[model]] = np.mean(time_buffer) if only_mean else time_buffer
    return timers


def get_plot(timers: Dict[str, Union[float, List[float]]]) -> plotly.graph_objs._figure.Figure:
    """
    Plot the inference time for each model.

    Parameters
    ----------
    timers : Dict[str, Union[float, List[float]]]
        Dictionary of inference times for each model.
    """
    data = pd.DataFrame.from_dict(timers, orient="columns")
    colors = ["#140f0d", "#2b2c4f", "#615aa2", "#a991fa"]
    fig = plotly.figure_factory.create_distplot(
        [data[col] for col in data.columns], data.columns, bin_size=0.2, colors=colors
    )
    fig.update_layout(title_text="Inference Time", xaxis_title="Inference Time (s)", yaxis_title="Number of Samples")
    return fig
        

st.set_page_config(page_title="Optimum Text Classification", page_icon="⭐")
st.title("⭐ Optimum Text Classification")
st.subheader("Classify financial news tone with 🤗 Optimum and ONNXRuntime")
st.markdown("""
[![GitHub](https://img.shields.io/badge/-%23121011.svg?style=for-the-badge&logo=github&logoColor=white)](https://github.com/ChainYo)
[![HuggingFace](https://img.shields.io/badge/-yellow.svg?style=for-the-badge&logo=data:image/svg+xml;base64,PD94bWwgdmVyc2lvbj0iMS4wIiBzdGFuZGFsb25lPSJubyI/Pgo8IURPQ1RZUEUgc3ZnIFBVQkxJQyAiLS8vVzNDLy9EVEQgU1ZHIDIwMDEwOTA0Ly9FTiIKICJodHRwOi8vd3d3LnczLm9yZy9UUi8yMDAxL1JFQy1TVkctMjAwMTA5MDQvRFREL3N2ZzEwLmR0ZCI+CjxzdmcgdmVyc2lvbj0iMS4wIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciCiB3aWR0aD0iMTc1LjAwMDAwMHB0IiBoZWlnaHQ9IjE3NS4wMDAwMDBwdCIgdmlld0JveD0iMCAwIDE3NS4wMDAwMDAgMTc1LjAwMDAwMCIKIHByZXNlcnZlQXNwZWN0UmF0aW89InhNaWRZTWlkIG1lZXQiPgoKPGcgdHJhbnNmb3JtPSJ0cmFuc2xhdGUoMC4wMDAwMDAsMTc1LjAwMDAwMCkgc2NhbGUoMC4xMDAwMDAsLTAuMTAwMDAwKSIKZmlsbD0iIzAwMDAwMCIgc3Ryb2tlPSJub25lIj4KPHBhdGggZD0iTTU2MyAxMjM2IGMtMjkgLTEzIC02MyAtNTkgLTYzIC04NiAwIC0yNiAzMyAtODAgNTIgLTg2IDE1IC00IDI2IDEKNDMgMjEgMjAgMjYgMjQgMjcgNTMgMTcgMjggLTkgMzMgLTggNDIgOCAxNyAzMiAxMSA2OSAtMTcgOTkgLTM0IDM3IC02OCA0NQotMTEwIDI3eiIvPgo8cGF0aCBkPSJNMTA2NCAxMjQwIGMtNTAgLTIwIC03NyAtODYgLTU0IC0xMzAgOSAtMTYgMTQgLTE3IDQyIC04IDI5IDEwIDMzIDkKNTUgLTE3IDIxIC0yNCAyNyAtMjYgNDggLTE3IDMxIDE0IDUxIDc2IDM2IDExNCAtMTcgNDYgLTg0IDc2IC0xMjcgNTh6Ii8+CjxwYXRoIGQ9Ik02MDAgODg4IGMwIC00OSAxNiAtOTggNTAgLTE1MSA4NSAtMTM0IDMyNSAtMTM0IDQxMCAwIDUxIDgwIDY5IDE4MwozMSAxODMgLTEwIDAgLTUwIC0xNSAtODcgLTMyIC02MCAtMjkgLTc5IC0zMyAtMTQ5IC0zMyAtNzAgMCAtODkgNCAtMTQ5IDMzCi0zNyAxNyAtNzcgMzIgLTg3IDMyIC0xNSAwIC0xOSAtNyAtMTkgLTMyeiIvPgo8L2c+Cjwvc3ZnPgo=)](https://huggingface.co/ChainYo)
[![LinkedIn](https://img.shields.io/badge/-%230077B5.svg?style=for-the-badge&logo=linkedin&logoColor=white)](https://www.linkedin.com/in/thomas-chaigneau-dev/)
[![Discord](https://img.shields.io/badge/Chainyo%233610-%237289DA.svg?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/)
""")

with st.expander("⭐ Details", expanded=True):
    st.markdown(
        """
        This app is a **demo** of the [🤗 Optimum Text Classification](https://huggingface.co/docs/optimum/onnxruntime/modeling_ort#optimum-inference-with-onnx-runtime) pipeline.
        We aim to compare the original pipeline with the ONNXRuntime pipeline.

        We use the [Finbert-Tone](https://huggingface.co/yiyanghkust/finbert-tone) model to classify financial news tone for the demo.
        
        You can enter multiple sentences to classify them by separating them with a `; (semicolon)`.
        """
    )

if "init_models" not in st.session_state:
    st.session_state["init_models"] = True
if st.session_state["init_models"]:
    with st.spinner(text="Loading files and models..."):
        loading_logs = st.empty()
        with loading_logs.container():
            BASE_PATH.mkdir(exist_ok=True)
            QUANTIZED_BASE_PATH.mkdir(exist_ok=True)
            OPTIMIZED_BASE_PATH.mkdir(exist_ok=True)

            if "tokenizer" not in st.session_state:
                tokenizer = BertTokenizer.from_pretrained(HUB_MODEL_PATH)
                st.session_state["tokenizer"] = tokenizer
            st.text("✅ Tokenizer loaded.")

            if "pt_model" not in st.session_state:
                pt_model = BertForSequenceClassification.from_pretrained(HUB_MODEL_PATH, num_labels=3)
                st.session_state["pt_model"] = pt_model
            st.text("✅ PyTorch model loaded.")

            if "ort_model" not in st.session_state:
                ort_model = ORTModelForSequenceClassification.from_pretrained(HUB_MODEL_PATH, from_transformers=True)
                # if not ONNX_MODEL_PATH.exists():
                #     ort_model.save_pretrained(ONNX_MODEL_PATH)
                st.session_state["ort_model"] = ort_model
            st.text("✅ ONNX Model loaded.")

            if "optimized_model" not in st.session_state:
                optimization_config = OptimizationConfig(optimization_level=99)
                optimizer = ORTOptimizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
                if not OPTIMIZED_MODEL_PATH.exists():
                    optimizer.export(ONNX_MODEL_PATH, OPTIMIZED_MODEL_PATH, optimization_config=optimization_config)
                    optimizer.model.config.save_pretrained(OPTIMIZED_BASE_PATH)
                optimized_model = ORTModelForSequenceClassification.from_pretrained(
                    OPTIMIZED_BASE_PATH, file_name=OPTIMIZED_MODEL_PATH.name
                )
                st.session_state["optimized_model"] = optimized_model
            st.text("✅ Optimized ONNX model loaded.")

            if "quantized_model" not in st.session_state:
                quantization_config = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
                quantizer = ORTQuantizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
                if not QUANTIZED_MODEL_PATH.exists():
                    quantizer.export(ONNX_MODEL_PATH, QUANTIZED_MODEL_PATH, quantization_config=quantization_config)
                    quantizer.model.config.save_pretrained(QUANTIZED_BASE_PATH)
                quantized_model = ORTModelForSequenceClassification.from_pretrained(
                    QUANTIZED_BASE_PATH, file_name=QUANTIZED_MODEL_PATH.name
                )
                st.session_state["quantized_model"] = quantized_model
            st.text("✅ Quantized ONNX model loaded.")

            if "pt_pipeline" not in st.session_state:
                pt_pipeline = pipeline(
                    "sentiment-analysis", tokenizer=st.session_state["tokenizer"], model=st.session_state["pt_model"]
                )
                st.session_state["pt_pipeline"] = pt_pipeline

            if "ort_pipeline" not in st.session_state:
                ort_pipeline = ort_pipeline(
                    "text-classification", tokenizer=st.session_state["tokenizer"], model=st.session_state["ort_model"]
                )
                st.session_state["ort_pipeline"] = ort_pipeline

            if "ort_optimized_pipeline" not in st.session_state:
                ort_optimized_pipeline = pipeline(
                    "text-classification", 
                    tokenizer=st.session_state["tokenizer"], 
                    model=st.session_state["optimized_model"],
                )
                st.session_state["ort_optimized_pipeline"] = ort_optimized_pipeline

            if "ort_quantized_pipeline" not in st.session_state:
                ort_quantized_pipeline = pipeline(
                    "text-classification", 
                    tokenizer=st.session_state["tokenizer"], 
                    model=st.session_state["quantized_model"],
                )
                st.session_state["ort_quantized_pipeline"] = ort_quantized_pipeline

            st.text("✅ All pipelines are ready.")
        sleep(2)
        loading_logs.success("🎉 Everything is ready!")
st.session_state["init_models"] = False

if "inference_timers" not in st.session_state:
    st.session_state["inference_timers"] = {}

exp_number = st.slider("The number of experiments per model.", min_value=100, max_value=300, value=150)
get_only_mean = st.checkbox("Get only the mean of the inference time for each model.", value=False)
input_text = st.text_area(
    "Enter text to classify", 
    "there is a shortage of capital, and we need extra financing; growth is strong and we have plenty of liquidity; there are doubts about our finances; profits are flat"
)
run_inference = st.button("🚀 Run inference")

if run_inference:
    st.text("🔎 Running inference...")
    sentences = input_text.split(";")
    st.session_state["inference_timers"] = get_timers(samples=sentences, exp_number=exp_number, only_mean=get_only_mean)
    st.plotly_chart(get_plot(st.session_state["inference_timers"]), use_container_width=True)