import gradio as gr

from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

tokenizer = AutoTokenizer.from_pretrained("coppercitylabs/uzbek-news-category-classifier")
model = AutoModelForSequenceClassification.from_pretrained("coppercitylabs/uzbek-news-category-classifier")

def prediction(news):
    # create pipeline
    clasifer = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model, return_all_scores=True)

    preds = clasifer(news)

    preds_dict={}
    for pred in preds[0]:
        preds_dict[pred['label']] = pred['score']

    return preds_dict


gradio_ui = gr.Interface(
    fn=prediction,
    title="O'zbek Yangiliklari Klassifikatsiyasi",
    description=f"",

    inputs=gr.inputs.Textbox(lines=10, label="Yangilik matnini kiriting"),
    outputs=gr.outputs.Label(num_top_classes=5, type="auto", label="Yangiliklar ko'rsatkichi"),
    theme="huggingface",
   )

gradio_ui.launch()