from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
import numpy as np
from scipy.special import softmax

MODEL = "Davlan/naija-twitter-sentiment-afriberta-large"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

# PT
model = AutoModelForSequenceClassification.from_pretrained(MODEL)


def get_senti(text):
    encoded_input = tokenizer(text, return_tensors='pt')
    output = model(**encoded_input)
    scores = output[0][0].detach().numpy()
    scores = softmax(scores)
    
    id2label = {0:"positive✅", 1:"neutral😐", 2:"negative❌"}
    
    ranking = np.argsort(scores)
    ranking = ranking[::-1]
    out = []
    for i in range(scores.shape[0]):
        l = id2label[ranking[i]]
        s = scores[ranking[i]]
        out.append(f"{i+1}) {l} {np.round(float(s), 4)}")
    
    out.append('\nNOTE: "The higher the values, the higher the sentiment for that class & vice-versa"')
    return "\n".join(out)
    



import gradio as gr


demo = gr.Interface(
    fn=get_senti,
    inputs=gr.Textbox(lines=2, placeholder="Enter Words Here (Igbo, Yoruba, Hausa Pidgin)..."),
    outputs="text",
)
demo.launch()