import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

MODEL_URL = "https://huggingface.co/dsfsi/PuoBERTa-News"
WEBSITE_URL = "https://www.kodiks.com/ai_solutions.html"

tokenizer = AutoTokenizer.from_pretrained("dsfsi/PuoBERTa-News")
model = AutoModelForSequenceClassification.from_pretrained("dsfsi/PuoBERTa-News")

categories = {
    "arts_culture_entertainment_and_media": "Botsweretshi, setso, boitapoloso le bobegakgang",
    "crime_law_and_justice": "Bosenyi, molao le bosiamisi",
    "disaster_accident_and_emergency_incident": "Masetlapelo, kotsi le tiragalo ya maemo a tshoganyetso",
    "economy_business_and_finance": "Ikonomi, tsa kgwebo le tsa ditšhelete",
    "education": "Thuto",
    "environment": "Tikologo",
    "health": "Boitekanelo",
    "politics": "Dipolotiki",
    "religion_and_belief": "Bodumedi le tumelo",
    "society": "Setšhaba"
}

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

    preds = clasifer(news)

    preds_dict = {}
    for pred in preds[0]:
        label = categories.get(pred['label'], pred['label'])
        preds_dict[label] = pred['score']

    return preds_dict

gradio_ui = gr.Interface(
    fn=prediction,
    title="Setswana News Classification",
    description=f"Enter Setswana news article to see the category of the news.\n For this classification, the {MODEL_URL} model was used.",
    examples=[
        ['Ka Letsatsi la Aforika, Aforika Borwa e tla be e keteka mabaka a boikemelo, le diketso tse di siameng tse e di dirileng go tokafatsa dikamano tsa yona le dinaga tse dingwe tsa Aforika.'],
        ["Thuto ya Setswana ke nngwe ya dithuto tse di botlhokwa mo sekolong se se tlhamaletseng go ruta bana ba ba mo lefatsheng la Botswana."],
        ["Mo kgweding e e fetileng, dipuisano tsa ditheko tsa dijalo di ile tsa tswelela, ka batho ba rekang le barui ba ba ruileng."],
        ["Masole a Aforika Borwa a ne a ya kwa Mozambique go tlisetsa motlakase morago ga maduo a kgatlha."],
    ],
    inputs=gr.inputs.Textbox(lines=10, label="Paste some Setswana news here"),
    outputs=gr.outputs.Label(num_top_classes=5, type="auto", label="News categories probabilities"),
    theme="huggingface",
    article="<p style='text-align: center'>For our other AI works: <a href='https://www.kodiks.com/ai_solutions.html' target='_blank'>https://www.kodiks.com/ai_solutions.html</a> | <a href='https://twitter.com/KodiksBilisim' target='_blank'>Contact us</a></p>",
)

gradio_ui.launch()