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 mengutuk POLRI atas kerusuhan yang terjadi di Malang",
    "Lesti mengatakan bahwa dia ingin mencabut gugatannya kepada Bilar di Kejaksaan Agung"
]

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()