from transformers import pipeline

import gradio as gr

pretrained_sentiment = "w11wo/indonesian-roberta-base-sentiment-classifier"
pretrained_ner = "cahya/bert-base-indonesian-NER"

sentiment_pipeline = pipeline(
    "sentiment-analysis",
    model=pretrained_sentiment,
    tokenizer=pretrained_sentiment,
    return_all_scores=True
)

ner_pipeline = pipeline(
    "ner",
    model=pretrained_ner,
    tokenizer=pretrained_ner
)

examples = [
    "Masyarakat sangat kecewa dengan tragedi Kanjuruhan",
    "Jokowi mengutuk kepolisian atas kerusuhan yang terjadi di Malang"
]

def sentiment_analysis(text):
    output = sentiment_pipeline(text)
    return {elm["label"]: elm["score"] for elm in output[0]}
    
def ner(text):
    output = ner_pipeline(text)
    return {"text": text, "entities": output}

demo = gr.Interface(
    fn=[sentiment_analysis, ner], 
    inputs=gr.Textbox(placeholder="Enter a sentence here..."), 
    outputs=["label", gr.HighlightedText()],
    interpretation=["default"],
    examples=[examples])

if __name__ == "__main__":
    demo.launch()