Update app.py

app.py

@@ -1,139 +1,131 @@
 import streamlit as st
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
 
-# Define a comprehensive dictionary of language names, codes, and model names
+# Define the dictionary of language models
 LANGUAGE_MODELS = {
-    'Afrikaans': ('af', 'Helsinki-NLP/opus-mt-en-af'),
-    'Albanian': ('sq', 'Helsinki-NLP/opus-mt-en-sq'),
-    'Arabic': ('ar', 'Helsinki-NLP/opus-mt-en-ar'),
-    'Armenian': ('hy', 'Helsinki-NLP/opus-mt-en-hy'),
-    'Catalan': ('ca', 'Helsinki-NLP/opus-mt-en-ca'),
-    'Czech': ('cs', 'Helsinki-NLP/opus-mt-en-cs'),
-    'Danish': ('da', 'Helsinki-NLP/opus-mt-en-da'),
-    'Dutch': ('nl', 'Helsinki-NLP/opus-mt-en-nl'),
-    'Esperanto': ('eo', 'Helsinki-NLP/opus-mt-en-eo'),
-    'Estonian': ('et', 'Helsinki-NLP/opus-mt-en-et'),
-    'Finnish': ('fi', 'Helsinki-NLP/opus-mt-en-fi'),
-    'French': ('fr', 'Helsinki-NLP/opus-mt-en-fr'),
-    'German': ('de', 'Helsinki-NLP/opus-mt-en-de'),
-    'Greek': ('el', 'Helsinki-NLP/opus-mt-en-el'),
-    'Gujarati': ('gu', 'Helsinki-NLP/opus-mt-en-gu'),
-    'Haitian Creole': ('ht', 'Helsinki-NLP/opus-mt-en-ht'),
-    'Hausa': ('ha', 'Helsinki-NLP/opus-mt-en-ha'),
-    'Hawaiian': ('haw', 'Helsinki-NLP/opus-mt-en-haw'),
-    'Hebrew': ('he', 'Helsinki-NLP/opus-mt-en-he'),
-    'Hindi': ('hi', 'Helsinki-NLP/opus-mt-en-hi'),
-    'Hungarian': ('hu', 'Helsinki-NLP/opus-mt-en-hu'),
-    'Icelandic': ('is', 'Helsinki-NLP/opus-mt-en-is'),
-    'Igbo': ('ig', 'Helsinki-NLP/opus-mt-en-ig'),
-    'Indonesian': ('id', 'Helsinki-NLP/opus-mt-en-id'),
-    'Irish': ('ga', 'Helsinki-NLP/opus-mt-en-ga'),
-    'Italian': ('it', 'Helsinki-NLP/opus-mt-en-it'),
-    'Japanese': ('ja', 'Helsinki-NLP/opus-mt-en-ja'),
-    'Javanese': ('jw', 'Helsinki-NLP/opus-mt-en-jw'),
-    'Kannada': ('kn', 'Helsinki-NLP/opus-mt-en-kn'),
-    'Khmer': ('km', 'Helsinki-NLP/opus-mt-en-km'),
-    'Korean': ('ko', 'Helsinki-NLP/opus-mt-en-ko'),
-    'Latin': ('la', 'Helsinki-NLP/opus-mt-en-la'),
-    'Latvian': ('lv', 'Helsinki-NLP/opus-mt-en-lv'),
-    'Lithuanian': ('lt', 'Helsinki-NLP/opus-mt-en-lt'),
-    'Luxembourgish': ('lb', 'Helsinki-NLP/opus-mt-en-lb'),
-    'Macedonian': ('mk', 'Helsinki-NLP/opus-mt-en-mk'),
-    'Malagasy': ('mg', 'Helsinki-NLP/opus-mt-en-mg'),
-    'Malayalam': ('ml', 'Helsinki-NLP/opus-mt-en-ml'),
-    'Maltese': ('mt', 'Helsinki-NLP/opus-mt-en-mt'),
-    'Maori': ('mi', 'Helsinki-NLP/opus-mt-en-mi'),
-    'Marathi': ('mr', 'Helsinki-NLP/opus-mt-en-mr'),
-    'Myanmar': ('my', 'Helsinki-NLP/opus-mt-en-my'),
-    'Nepali': ('ne', 'Helsinki-NLP/opus-mt-en-ne'),
-    'Norwegian': ('no', 'Helsinki-NLP/opus-mt-en-no'),
-    'Nyanja': ('ny', 'Helsinki-NLP/opus-mt-en-ny'),
-    'Odia': ('or', 'Helsinki-NLP/opus-mt-en-or'),
-    'Oromo': ('om', 'Helsinki-NLP/opus-mt-en-om'),
-    'Pashto': ('ps', 'Helsinki-NLP/opus-mt-en-ps'),
-    'Persian': ('fa', 'Helsinki-NLP/opus-mt-en-fa'),
-    'Polish': ('pl', 'Helsinki-NLP/opus-mt-en-pl'),
-    'Portuguese': ('pt', 'Helsinki-NLP/opus-mt-en-pt'),
-    'Punjabi': ('pa', 'Helsinki-NLP/opus-mt-en-pa'),
-    'Quechua': ('qu', 'Helsinki-NLP/opus-mt-en-qu'),
-    'Romanian': ('ro', 'Helsinki-NLP/opus-mt-en-ro'),
-    'Russian': ('ru', 'Helsinki-NLP/opus-mt-en-ru'),
-    'Samoan': ('sm', 'Helsinki-NLP/opus-mt-en-sm'),
-    'Scots Gaelic': ('gd', 'Helsinki-NLP/opus-mt-en-gd'),
-    'Serbian': ('sr', 'Helsinki-NLP/opus-mt-en-sr'),
-    'Sesotho': ('st', 'Helsinki-NLP/opus-mt-en-st'),
-    'Shona': ('sn', 'Helsinki-NLP/opus-mt-en-sn'),
-    'Sindhi': ('sd', 'Helsinki-NLP/opus-mt-en-sd'),
-    'Sinhala': ('si', 'Helsinki-NLP/opus-mt-en-si'),
-    'Slovak': ('sk', 'Helsinki-NLP/opus-mt-en-sk'),
-    'Slovenian': ('sl', 'Helsinki-NLP/opus-mt-en-sl'),
-    'Somali': ('so', 'Helsinki-NLP/opus-mt-en-so'),
-    'Spanish': ('es', 'Helsinki-NLP/opus-mt-en-es'),
-    'Sundanese': ('su', 'Helsinki-NLP/opus-mt-en-su'),
-    'Swahili': ('sw', 'Helsinki-NLP/opus-mt-en-sw'),
-    'Swedish': ('sv', 'Helsinki-NLP/opus-mt-en-sv'),
-    'Tajik': ('tg', 'Helsinki-NLP/opus-mt-en-tg'),
-    'Tamil': ('ta', 'Helsinki-NLP/opus-mt-en-ta'),
-    'Telugu': ('te', 'Helsinki-NLP/opus-mt-en-te'),
-    'Thai': ('th', 'Helsinki-NLP/opus-mt-en-th'),
-    'Turkmen': ('tk', 'Helsinki-NLP/opus-mt-en-tk'),
-    'Turkish': ('tr', 'Helsinki-NLP/opus-mt-en-tr'),
-    'Ukrainian': ('uk', 'Helsinki-NLP/opus-mt-en-uk'),
-    'Urdu': ('ur', 'Helsinki-NLP/opus-mt-en-ur'),
-    'Vietnamese': ('vi', 'Helsinki-NLP/opus-mt-en-vi'),
-    'Welsh': ('cy', 'Helsinki-NLP/opus-mt-en-cy'),
-    'Xhosa': ('xh', 'Helsinki-NLP/opus-mt-en-xh'),
-    'Yiddish': ('yi', 'Helsinki-NLP/opus-mt-en-yi'),
-    'Yoruba': ('yo', 'Helsinki-NLP/opus-mt-en-yo'),
-    'Zulu': ('zu', 'Helsinki-NLP/opus-mt-en-zu'),
+    'Afrikaans': 'af',
+    'Albanian': 'sq',
+    'Amharic': 'am',
+    'Arabic': 'ar',
+    'Armenian': 'hy',
+    'Bengali': 'bn',
+    'Bosnian': 'bs',
+    'Catalan': 'ca',
+    'Croatian': 'hr',
+    'Czech': 'cs',
+    'Danish': 'da',
+    'Dutch': 'nl',
+    'Esperanto': 'eo',
+    'Estonian': 'et',
+    'Finnish': 'fi',
+    'French': 'fr',
+    'German': 'de',
+    'Greek': 'el',
+    'Gujarati': 'gu',
+    'Haitian Creole': 'ht',
+    'Hausa': 'ha',
+    'Hawaiian': 'haw',
+    'Hebrew': 'he',
+    'Hindi': 'hi',
+    'Hungarian': 'hu',
+    'Icelandic': 'is',
+    'Igbo': 'ig',
+    'Indonesian': 'id',
+    'Irish': 'ga',
+    'Italian': 'it',
+    'Japanese': 'ja',
+    'Javanese': 'jw',
+    'Kannada': 'kn',
+    'Khmer': 'km',
+    'Korean': 'ko',
+    'Latin': 'la',
+    'Latvian': 'lv',
+    'Lithuanian': 'lt',
+    'Luxembourgish': 'lb',
+    'Macedonian': 'mk',
+    'Malagasy': 'mg',
+    'Malayalam': 'ml',
+    'Maltese': 'mt',
+    'Maori': 'mi',
+    'Marathi': 'mr',
+    'Myanmar': 'my',
+    'Nepali': 'ne',
+    'Norwegian': 'no',
+    'Nyanja': 'ny',
+    'Odia': 'or',
+    'Oromo': 'om',
+    'Pashto': 'ps',
+    'Persian': 'fa',
+    'Polish': 'pl',
+    'Portuguese': 'pt',
+    'Punjabi': 'pa',
+    'Quechua': 'qu',
+    'Romanian': 'ro',
+    'Russian': 'ru',
+    'Samoan': 'sm',
+    'Scots Gaelic': 'gd',
+    'Serbian': 'sr',
+    'Sesotho': 'st',
+    'Shona': 'sn',
+    'Sindhi': 'sd',
+    'Sinhala': 'si',
+    'Slovak': 'sk',
+    'Slovenian': 'sl',
+    'Somali': 'so',
+    'Spanish': 'es',
+    'Sundanese': 'su',
+    'Swahili': 'sw',
+    'Swedish': 'sv',
+    'Tajik': 'tg',
+    'Tamil': 'ta',
+    'Telugu': 'te',
+    'Thai': 'th',
+    'Turkmen': 'tk',
+    'Turkish': 'tr',
+    'Ukrainian': 'uk',
+    'Urdu': 'ur',
+    'Vietnamese': 'vi',
+    'Welsh': 'cy',
+    'Xhosa': 'xh',
+    'Yiddish': 'yi',
+    'Yoruba': 'yo',
+    'Zulu': 'zu',
 }
 
 @st.cache_resource
-def load_model(target_language):
-    code, model_name = LANGUAGE_MODELS.get(target_language, (None, None))
-    if not model_name:
-        st.error(f"Model for language '{target_language}' not found.")
-        return None, None
-    
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
+def load_model():
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
     return tokenizer, model
 
-def translate_text(text, target_language):
-    tokenizer, model = load_model(target_language)
-    if tokenizer is None or model is None:
-        return ""
-    
-    # Tokenize the input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True)
+def translate(text, target_language):
+    tokenizer, model = load_model()
     
-    # Generate translation
-    translated = model.generate(**inputs)
-    
-    # Decode the translated text
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-    return translated_text
+    # Set the target language code for translation
+    target_lang_code = LANGUAGE_MODELS.get(target_language)
+    if not target_lang_code:
+        return "Target language not supported."
 
-def main():
-    st.title("Language Translator")
-    st.write("Translate English Text into Your Preferred Language")
+    tokenizer.src_lang = "en"
+    encoded_input = tokenizer(text, return_tensors="pt")
+    generated_tokens = model.generate(**encoded_input, forced_bos_token_id=tokenizer.get_lang_id(target_lang_code))
+    translation = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+    return translation
 
-    # Input text from the user
-    source_text = st.text_area("Enter text in English:", "")
-    
-    # Select target language
-    target_language = st.selectbox(
-        "Select target language:",
-        options=list(LANGUAGE_MODELS.keys())
-    )
-    
-    if st.button("Translate"):
-        if source_text:
-            translated_text = translate_text(source_text, target_language)
-            st.write(f"Translated text ({target_language}):")
-            st.write(translated_text)
-        else:
-            st.warning("Please enter text to translate.")
+st.title('English to Any Language Translator')
+
+text_input = st.text_area("Enter text in English:", "Hello, how are you?")
+
+target_language = st.selectbox(
+    'Select the target language:',
+    list(LANGUAGE_MODELS.keys())
+)
 
-if __name__ == "__main__":
-    main()
+if st.button('Translate'):
+    with st.spinner('Translating...'):
+        try:
+            translation = translate(text_input, target_language)
+            st.write(f'Translation ({target_language}):')
+            st.write(translation)
+        except Exception as e:
+            st.error(f"Error: {e}")