Update app.py

app.py

@@ -1,12 +1,10 @@
-from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, pipeline
+from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, MarianMTModel, MarianTokenizer, pipeline
 import nltk.data
 import pandas as pd
 import matplotlib.pyplot as plt
-from google_trans_new import google_translator  
 
 nltk.download('punkt')
 
-
 import gradio as gr
 from gradio.mix import Parallel
 
@@ -18,6 +16,12 @@ pretrained_ner = "cahya/bert-base-indonesian-NER"
 
 sentence_tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
 
+tokenizer_translate = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-id-en")
+model_translate = MarianMTModel.from_pretrained(
+    "Helsinki-NLP/opus-mt-id-en")
+#finetuned_model = MarianMTModel.from_pretrained(
+#    "wolfrage89/annual_report_translation_id_en")
+
 sentiment_pipeline = pipeline(
     "sentiment-analysis",
     model=pretrained_sentiment,
@@ -32,14 +36,19 @@ ner_pipeline = pipeline(
     grouped_entities=True
 )
 
-translator = google_translator()
-
 examples = [
     "Perusahaan industri e-commerce Indonesia, Bukalapak telah memberhentikan puluhan karyawan dari beberapa function; Berlawanan dengan PHK sebelumnya, perusahaan mengontrak jajaran pekerja kantornya, harian Kompas melaporkan.",
     "Dengan pabrik produksi baru, perusahaan akan meningkatkan kapasitasnya untuk memenuhi peningkatan permintaan yang diharapkan dan akan meningkatkan penggunaan bahan baku dan oleh karena itu meningkatkan profitabilitas produksi.",
     "Lifetree didirikan pada tahun 2000, dan pendapatannya meningkat rata-rata 40% dengan margin di akhir 30-an."
 ]
 
+def get_translation(text):
+    translated_tokens = model_translate.generate(
+        **tokenizer_translate([text], return_tensors='pt', max_length=104, truncation=True))[0]
+    translated_sentence = tokenizer_translate.decode(
+        translated_tokens, skip_special_tokens=True)
+    return translated_sentence
+
 def summ_t5(text):
     input_ids = tokenizer_t5.encode(text, return_tensors='pt')
     summary_ids = model_t5.generate(input_ids,
@@ -66,7 +75,7 @@ def ner(text):
 def sentiment_df(text):
     df = pd.DataFrame(columns=['Text', 'Eng', 'Label', 'Score'])
     text_list = sentence_tokenizer.tokenize(text)
-    eng_text = [translator.translate(text, lang_src='id',lang_tgt='en') for text in text_list]
+    eng_text = [get_translation(text) for text in text_list]
     result = [sentiment_analysis(text) for text in text_list]
     labels = []
     scores = []
@@ -82,7 +91,7 @@ def sentiment_df(text):
     
 def run(text):
     summ_ = summ_t5(text)
-    summ_translated = translator.translate(summ_, lang_src='id',lang_tgt='en')
+    summ_translated = get_translation(summ_)
     sent_ = sentiment_analysis(summ_translated )
     ner_ = ner(summ_)
     df_sentiment = sentiment_df(text)