Update app.py

app.py

@@ -1,4 +1,6 @@
-from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, BertTokenizer, EncoderDecoderModel, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, pipeline
+import nltk.data
+
 
 import gradio as gr
 from gradio.mix import Parallel
@@ -6,13 +8,24 @@ from gradio.mix import Parallel
 tokenizer_t5 = T5Tokenizer.from_pretrained("panggi/t5-base-indonesian-summarization-cased")
 model_t5 = T5ForConditionalGeneration.from_pretrained("panggi/t5-base-indonesian-summarization-cased")
 
-tokenizer_bert = BertTokenizer.from_pretrained("cahya/bert2bert-indonesian-summarization")
-tokenizer_bert.bos_token = tokenizer_bert.cls_token
-tokenizer_bert.eos_token = tokenizer_bert.sep_token
-model_bert = EncoderDecoderModel.from_pretrained("cahya/bert2bert-indonesian-summarization")
+pretrained_sentiment = "w11wo/indonesian-roberta-base-sentiment-classifier"
+pretrained_ner = "cahya/bert-base-indonesian-NER"
+
+sentence_tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
 
-t5_para_tokenizer = AutoTokenizer.from_pretrained("Wikidepia/IndoT5-base-paraphrase")  
-t5_para_model = AutoModelForSeq2SeqLM.from_pretrained("Wikidepia/IndoT5-base-paraphrase")
+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,
+    grouped_entities=True
+)
 
 
 def summ_t5(text):
@@ -28,45 +41,44 @@ def summ_t5(text):
     summary_text = tokenizer_t5.decode(summary_ids[0], skip_special_tokens=True)
     return summary_text
     
-def summ_bert(text):
-    input_ids = tokenizer_bert.encode(text, return_tensors="pt")
-    summary_ids= model_bert.generate(input_ids,
-                max_length=100,
-                num_beams=10,
-                repetition_penalty=2.5,
-                length_penalty=1.0,
-                early_stopping=True,
-                no_repeat_ngram_size=2,
-                use_cache=True)
-
-    summary_text = tokenizer_bert.decode(summary_ids[0], skip_special_tokens=True)
-    return summary_text
+def sentiment_analysis(text):
+    output = sentiment_pipeline(text)
+    return {elm["label"]: elm["score"] for elm in output[0]}
 
-def para_t5(text):
-    encoding = t5_para_tokenizer(text, padding='longest', return_tensors='pt')
-    outputs = t5_para_model.generate(
-                input_ids=encoding["input_ids"],
-                attention_mask=encoding["attention_mask"],
-                max_length=100,
-                do_sample=True,
-                top_k=120,
-                top_p=0.95,
-                early_stopping=True,
-                num_return_sequences=1)
-    return t5_para_tokenizer.decode(
-            outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True
-        )
+def ner(text):
+    output = ner_pipeline(text)
+    for elm in output:
+        elm['entity'] = elm['entity_group']
+    return {"text": text, "entities": output}
     
-
-def summarize(text):
-    t5_ = summ_t5(text)
-    bert_ = summ_bert(text)
-    para_ = para_t5(t5_)
-    return t5_, bert_, para_
+def sentiment_df(text):
+    text_list = tokenizer.tokenize(text)
+    result = [sentiment_analysis(text) for text in text_list]
+    sentence = []
+    labels = []
+    scores = []
+    for pred in result:
+        idx = list(pred.values()).index(max(list(pred.values())))
+        labels.append(list(pred.keys())[idx])
+        scores.append(round(list(pred.values())[idx], 3))
+    df['Text'] = text_list
+    df['Label'] = labels
+    df['Score'] = scores
+    return df
+    
+def run(text):
+    summ_ = summ_t5(text)
+    sent_ = sentiment_analysis(summ_)
+    ner_ = ner(summ_)
+    df_ = sentiment_df(text)
+    ner_all = ner(text)
+    fig = plt.figure()
+    df.groupby(["Label"])["Text"].count().plot.pie(autopct="%.1f%%", figsize=(6,6))
+    return summ_, sent_, ner_, fig, ner_all
 
 if __name__ == "__main__":
     with gr.Blocks() as demo:
-        gr.Markdown("""<h1 style="text-align:center">Summary of Summarizer - Indonesia</h1>""")
+        gr.Markdown("""<h1 style="text-align:center">News Analyzer - Indonesia</h1>""")
 
         gr.Markdown(
             """
@@ -77,9 +89,14 @@ if __name__ == "__main__":
             with gr.Column():
                 input_text = gr.Textbox(label="Input Text")
                 analyze_button = gr.Button(label="Analyze")
+                
+                summ_output = gr.Textbox(label="Article Summary")
+                ner_output = gr.HighlightedText(label="NER Summary")
+                sent_output = gr.Textbox(label="Sentiment Summary")
+                
             with gr.Column():
-                t5_output = gr.Textbox(label="T5 Base Output")
-                bert_output = gr.Textbox(label="Bert2Bert Base Output")
-                para_output = gr.Textbox(label="T5 Paraphrase Output")
-        analyze_button.click(summarize, inputs=input_text, outputs=[t5_output, bert_output, para_output])
+                plot_component = gr.Plot(label="Pie Chart of Sentiments")
+                ner_all_output = gr.HighlightedText(label="NER Article")
+                
+        analyze_button.click(run, inputs=input_text, outputs=[summ_output, sent_output, ner_output, plot_component, ner_all_output])
     demo.launch()