Update app.py

app.py

@@ -1,65 +1,70 @@
 import gradio as gr
-from googletrans import Translator
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-european_languages = {
-    'afrikaans': 'af',
-    'albanian': 'sq',
-    'armenian': 'hy',
-    'azerbaijani': 'az',
-    'basque': 'eu',
-    'belarusian': 'be',
-    'bosnian': 'bs',
-    'bulgarian': 'bg',
-    'catalan': 'ca',
-    'croatian': 'hr',
-    'czech': 'cs',
-    'danish': 'da',
-    'dutch': 'nl',
-    'english': 'en',
-    'estonian': 'et',
-    'finnish': 'fi',
-    'french': 'fr',
-    'georgian': 'ka',
-    'german': 'de',
-    'greek': 'el',
-    'hungarian': 'hu',
-    'icelandic': 'is',
-    'irish': 'ga',
-    'italian': 'it',
-    'latvian': 'lv',
-    'lithuanian': 'lt',
-    'luxembourgish': 'lb',
-    'macedonian': 'mk',
-    'maltese': 'mt',
-    'norwegian': 'no',
-    'polish': 'pl',
-    'portuguese': 'pt',
-    'romanian': 'ro',
-    'russian': 'ru',
-    'scots gaelic': 'gd',
-    'serbian': 'sr',
-    'slovak': 'sk',
-    'slovenian': 'sl',
-    'spanish': 'es',
-    'swedish': 'sv',
-    'ukrainian': 'uk',
-    'welsh': 'cy',
+languages = {
+    'Afrikaans': 'afr_Latn',
+    'Albanian': 'als_Latn',
+    'Basque': 'eus_Latn',
+    'Belarusian': 'bel_Cyrl',
+    'Bosnian': 'bos_Latn',
+    'Bulgarian': 'bul_Cyrl',
+    'Catalan': 'cat_Latn',
+    'Croatian': 'hrv_Latn',
+    'Czech': 'ces_Latn',
+    'Danish': 'dan_Latn',
+    'Dutch': 'nld_Latn',
+    'English': 'eng_Latn',
+    'Estonian': 'est_Latn',
+    'Finnish': 'fin_Latn',
+    'French': 'fra_Latn',
+    'Galician': 'glg_Latn',
+    'German': 'deu_Latn',
+    'Greek': 'ell_Grek',
+    'Hungarian': 'hun_Latn',
+    'Icelandic': 'isl_Latn',
+    'Irish': 'gle_Latn',
+    'Italian': 'ita_Latn',
+    'Lithuanian': 'lit_Latn',
+    'Luxembourgish': 'ltz_Latn',
+    'Macedonian': 'mkd_Cyrl',
+    'Maltese': 'mlt_Latn',
+    'Norwegian Bokmål': 'nob_Latn',
+    'Norwegian Nynorsk': 'nno_Latn',
+    'Polish': 'pol_Latn',
+    'Portuguese': 'por_Latn',
+    'Romanian': 'ron_Latn',
+    'Russian': 'rus_Cyrl',
+    'Serbian': 'srp_Cyrl',
+    'Slovak': 'slk_Latn',
+    'Slovenian': 'slv_Latn',
+    'Spanish': 'spa_Latn',
+    'Swedish': 'swe_Latn',
+    'Ukrainian': 'ukr_Cyrl',
+    'Welsh': 'cym_Latn'
 }
 
-translator = Translator()
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
 
 def translate_article(article, language):
-    lang = translator.detect(article).lang
-    lang_code = european_languages[language]
-    translated_text = translator.translate(text, src=lang, dest= lang_code)
-    return translated_text.text
+    inputs = tokenizer(article, return_tensors="pt")
+    if language=='English':
+      lang_code='eng_Latn'
+      print("Yes")
+    else:
+      lang_code = languages[language]
+      print("No")
+    print(lang_code)
+    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code])
+    result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+    return result
 
-language_choices = list(european_languages.keys())
+language_choices = list(languages.keys())
 
 iface = gr.Interface(
     fn=translate_article,
     inputs=["text",gr.Dropdown(
-            language_choices, value="english", multiselect=False, label="Choose the language.")],
+            language_choices, value="English", multiselect=False, label="Choose the language.")],
     outputs="text",
     title="Translation Tool"
 )