import gradio as gr
import time
from optimum.intel.openvino import OVModelForSequenceClassification
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

model_id = "juliensimon/xlm-v-base-language-id"

p = pipeline("text-classification", model=model_id)

ov_model = OVModelForSequenceClassification.from_pretrained(
    model_id, from_transformers=True
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

ov_p = pipeline("text-classification", model=ov_model, tokenizer=tokenizer)

# Warmum
for i in range(100):
    ov_p("Hello world")


def process(text, model, top_k=5):
    if model == "Vanilla":
        pipe = p
    else:
        pipe = ov_p
    tick = time.time()
    pred = pipe(text, top_k=top_k)
    tock = time.time()
    scores = {x["label"]: x["score"] for x in pred}
    msg = f"{(tock-tick)*1000:.2f} milliseconds"
    return scores, msg


# Gradio inputs
input_text = gr.Text(label="Enter text")
input_model = gr.Radio(label="Pick a model", choices=["Vanilla", "OpenVINO"])

# Gradio outputs
output_labels = gr.Label(label="Languages", num_top_classes=5)
output_text = gr.Text(label="Prediction time")

description = "This Space lets you perform language identification on the 102 languages present in the google/fleurs dataset. The underlying model scores 99.3% accuracy on the validation set. Inference is optimized with Optimum Intel and OpenVINO."

iface = gr.Interface(
    description=description,
    fn=process,
    inputs=[input_text, input_model],
    outputs=[output_labels, output_text],
    examples=[
        [
            "Kila mtu ana haki ya kuelimishwa. Elimu yapasa itolewe bure hasa ile ya madarasa ya chini. Elimu ya masarasa ya chini ihudhuriwe kwa lazima. Elimu ya ufundi na ustadi iwe wazi kwa wote. Na elimu ya juu iwe wazi kwa wote kwa kutegemea sifa ya mtu",
            "OpenVINO",
        ],
        [
            "Ang bawat tao'y may karapatan sa edukasyon. Ang edukasyon ay walang bayad, doon man lamang sa elementarya at pangunahing antas. Ang edukasyong elementarya ay magiging sapilitan. Ang edukasyong teknikal at propesyonal ay gagawing maabot ng lahat at ang lalong mataas na edukasyon ay ipagkakaloob nang pantay-pantay sa lahat batay sa pagiging karapat-dapat.",
            "OpenVINO",
        ],
        [
            "Každý má právo na vzdělání. Vzdělání nechť je bezplatné, alespoň v počátečních a základních stupních. Základní vzdělání je povinné. Technické a odborné vzdělání budiž všeobecně přístupné a rovněž vyšší vzdělání má být stejně přístupné všem podle schopností.",
            "OpenVINO",
        ],
        [
            "Baat yii laa mën a wax ci jéego yi Marubba Faal séqi jëm ci bindkat yi taamu làmmiñi réew mi. Am naa ci mbégte lool, ngir yàgg naa ko wax ne bumu xeeb li gafakaam dencal xam-xam yu bawoo ca Maam. Téere bii mu tudde Yóbbalu Ndaw firndeel na sama njort.",
            "OpenVINO",
        ],
    ],
    allow_flagging="never",
)

iface.launch()