"""⭐ Text Classification with Optimum and ONNXRuntime

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

import streamlit as st

from transformers import AutoTokenizer, AutoModel, pipeline
from optimum.onnxruntime import ORTModelForSequenceClassification
from optimum.pipelines import pipeline


MODEL_PATH = "cardiffnlp/twitter-roberta-base-sentiment-latest"

st.set_page_config(page_title="Optimum Text Classification", page_icon="⭐")
st.title("🤗 Optimum Text Classification")
st.subheader("Sentiment analysis 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/)
""")

if "tokenizer" not in st.session_state:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    st.session_state["tokenizer"] = tokenizer

if "ort_model" not in st.session_state:
    ort_model = ORTModelForSequenceClassification.from_pretrained(MODEL_PATH, from_transformers=True)
    st.session_state["ort_model"] = ort_model

if "pt_model" not in st.session_state:
    pt_model = AutoModel.from_pretrained(MODEL_PATH)
    st.session_state["pt_model"] = pt_model

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

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


model_format = st.radio("Choose the model format", ("PyTorch", "ONNXRuntime"))
optimized = st.checkbox("Optimize the model for inference", value=False)
quantized = st.checkbox("Quantize the model", value=False)

if model_format == "PyTorch":
    optimized.disabled = True
    quantized.disabled = True