from transformers import pipeline

import gradio as gr
from gradio.mix import Parallel

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 = [
    "Jokowi sangat kecewa dengan POLRI atas kerusuhan yang terjadi di Malang",
    "Lesti marah terhadap perlakuan KDRT yang dilakukan oleh Bilar",
    "Ungkapan rasa bahagia diutarakan oleh Coki Pardede karena kebabasannya dari penjara"
]

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}

sentiment_demo = gr.Interface(
    fn=sentiment_analysis,
    inputs="text",
    outputs="label")

ner_demo = gr.Interface(
    ner,
    "text",
    gr.HighlightedText(),
    examples=examples)

if __name__ == "__main__":
    Parallel(sentiment_demo, ner_demo,
    inputs=gr.Textbox(lines=10, label="Input Text", placeholder="Enter sentences here..."),
    title="Entity Based Sentiment Analysis Indonesia",
    examples=examples).launch()