Update app.py

app.py

@@ -1,132 +1,52 @@
 import streamlit as st
 import os
-import io
-from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
 import time
-import json
-from typing import List
 import torch
-import random
 import logging
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+
+# Set Streamlit page configuration
+st.set_page_config(page_title="M2M100 Translator")
 
+# Check device
 if torch.cuda.is_available():
     device = torch.device("cuda:0")
 else:
     device = torch.device("cpu")
     logging.warning("GPU not found, using CPU, translation will be very slow.")
 
-st.cache(suppress_st_warning=True, allow_output_mutation=True)
-st.set_page_config(page_title="M2M100 Translator")
-
+# Language code mapping
 lang_id = {
-    "Afrikaans": "af",
-    "Amharic": "am",
-    "Arabic": "ar",
-    "Asturian": "ast",
-    "Azerbaijani": "az",
-    "Bashkir": "ba",
-    "Belarusian": "be",
-    "Bulgarian": "bg",
-    "Bengali": "bn",
-    "Breton": "br",
-    "Bosnian": "bs",
-    "Catalan": "ca",
-    "Cebuano": "ceb",
-    "Czech": "cs",
-    "Welsh": "cy",
-    "Danish": "da",
-    "German": "de",
-    "Greeek": "el",
-    "English": "en",
-    "Spanish": "es",
-    "Estonian": "et",
-    "Persian": "fa",
-    "Fulah": "ff",
-    "Finnish": "fi",
-    "French": "fr",
-    "Western Frisian": "fy",
-    "Irish": "ga",
-    "Gaelic": "gd",
-    "Galician": "gl",
-    "Gujarati": "gu",
-    "Hausa": "ha",
-    "Hebrew": "he",
-    "Hindi": "hi",
-    "Croatian": "hr",
-    "Haitian": "ht",
-    "Hungarian": "hu",
-    "Armenian": "hy",
-    "Indonesian": "id",
-    "Igbo": "ig",
-    "Iloko": "ilo",
-    "Icelandic": "is",
-    "Italian": "it",
-    "Japanese": "ja",
-    "Javanese": "jv",
-    "Georgian": "ka",
-    "Kazakh": "kk",
-    "Central Khmer": "km",
-    "Kannada": "kn",
-    "Korean": "ko",
-    "Luxembourgish": "lb",
-    "Ganda": "lg",
-    "Lingala": "ln",
-    "Lao": "lo",
-    "Lithuanian": "lt",
-    "Latvian": "lv",
-    "Malagasy": "mg",
-    "Macedonian": "mk",
-    "Malayalam": "ml",
-    "Mongolian": "mn",
-    "Marathi": "mr",
-    "Malay": "ms",
-    "Burmese": "my",
-    "Nepali": "ne",
-    "Dutch": "nl",
-    "Norwegian": "no",
-    "Northern Sotho": "ns",
-    "Occitan": "oc",
-    "Oriya": "or",
-    "Panjabi": "pa",
-    "Polish": "pl",
-    "Pushto": "ps",
-    "Portuguese": "pt",
-    "Romanian": "ro",
-    "Russian": "ru",
-    "Sindhi": "sd",
-    "Sinhala": "si",
-    "Slovak": "sk",
-    "Slovenian": "sl",
-    "Somali": "so",
-    "Albanian": "sq",
-    "Serbian": "sr",
-    "Swati": "ss",
-    "Sundanese": "su",
-    "Swedish": "sv",
-    "Swahili": "sw",
-    "Tamil": "ta",
-    "Thai": "th",
-    "Tagalog": "tl",
-    "Tswana": "tn",
-    "Turkish": "tr",
-    "Ukrainian": "uk",
-    "Urdu": "ur",
-    "Uzbek": "uz",
-    "Vietnamese": "vi",
-    "Wolof": "wo",
-    "Xhosa": "xh",
-    "Yiddish": "yi",
-    "Yoruba": "yo",
-    "Chinese": "zh",
-    "Zulu": "zu",
+    "Afrikaans": "af", "Amharic": "am", "Arabic": "ar", "Asturian": "ast",
+    "Azerbaijani": "az", "Bashkir": "ba", "Belarusian": "be", "Bulgarian": "bg",
+    "Bengali": "bn", "Breton": "br", "Bosnian": "bs", "Catalan": "ca",
+    "Cebuano": "ceb", "Czech": "cs", "Welsh": "cy", "Danish": "da",
+    "German": "de", "Greeek": "el", "English": "en", "Spanish": "es",
+    "Estonian": "et", "Persian": "fa", "Fulah": "ff", "Finnish": "fi",
+    "French": "fr", "Western Frisian": "fy", "Irish": "ga", "Gaelic": "gd",
+    "Galician": "gl", "Gujarati": "gu", "Hausa": "ha", "Hebrew": "he",
+    "Hindi": "hi", "Croatian": "hr", "Haitian": "ht", "Hungarian": "hu",
+    "Armenian": "hy", "Indonesian": "id", "Igbo": "ig", "Iloko": "ilo",
+    "Icelandic": "is", "Italian": "it", "Japanese": "ja", "Javanese": "jv",
+    "Georgian": "ka", "Kazakh": "kk", "Central Khmer": "km", "Kannada": "kn",
+    "Korean": "ko", "Luxembourgish": "lb", "Ganda": "lg", "Lingala": "ln",
+    "Lao": "lo", "Lithuanian": "lt", "Latvian": "lv", "Malagasy": "mg",
+    "Macedonian": "mk", "Malayalam": "ml", "Mongolian": "mn", "Marathi": "mr",
+    "Malay": "ms", "Burmese": "my", "Nepali": "ne", "Dutch": "nl",
+    "Norwegian": "no", "Northern Sotho": "ns", "Occitan": "oc", "Oriya": "or",
+    "Panjabi": "pa", "Polish": "pl", "Pushto": "ps", "Portuguese": "pt",
+    "Romanian": "ro", "Russian": "ru", "Sindhi": "sd", "Sinhala": "si",
+    "Slovak": "sk", "Slovenian": "sl", "Somali": "so", "Albanian": "sq",
+    "Serbian": "sr", "Swati": "ss", "Sundanese": "su", "Swedish": "sv",
+    "Swahili": "sw", "Tamil": "ta", "Thai": "th", "Tagalog": "tl",
+    "Tswana": "tn", "Turkish": "tr", "Ukrainian": "uk", "Urdu": "ur",
+    "Uzbek": "uz", "Vietnamese": "vi", "Wolof": "wo", "Xhosa": "xh",
+    "Yiddish": "yi", "Yoruba": "yo", "Chinese": "zh", "Zulu": "zu",
 }
 
-
-@st.cache_resource(suppress_st_warning=True, allow_output_mutation=True)
-def load_model(
-    pretrained_model: str = "facebook/m2m100_1.2B",
-    cache_dir: str = "models/",
-):
+# Cache the model and tokenizer using new API
+@st.cache_resource
+def load_model(pretrained_model="facebook/m2m100_1.2B", cache_dir="models/"):
     tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
     model = M2M100ForConditionalGeneration.from_pretrained(
         pretrained_model, cache_dir=cache_dir
@@ -134,40 +54,48 @@ def load_model(
     model.eval()
     return tokenizer, model
 
+# App Title and Intro
+st.title("🌐 M2M100 Translator")
+st.write("""
+M2M100 is a multilingual encoder-decoder (seq-to-seq) model trained for Many-to-Many multilingual translation.
+It supports **100 languages** and translates in **9900 directions**.  
+Model: `facebook/m2m100_1.2B`  
+More info: [Paper](https://arxiv.org/abs/2010.11125) | [Repo](https://github.com/pytorch/fairseq/tree/master/examples/m2m_100)
+""")
 
-st.title("M2M100 Translator")
-st.write("M2M100 is a multilingual encoder-decoder (seq-to-seq) model trained for Many-to-Many multilingual translation. It was introduced in this paper https://arxiv.org/abs/2010.11125 and first released in https://github.com/pytorch/fairseq/tree/master/examples/m2m_100 repository. The model that can directly translate between the 9,900 directions of 100 languages.\n")
-
-st.write(" This demo uses the facebook/m2m100_1.2B model. For local inference see https://github.com/ikergarcia1996/Easy-Translate")
-
-
-user_input: str = st.text_area(
-    "Input text",
+# Input Text Area
+user_input = st.text_area(
+    "Enter text to translate:",
     height=200,
     max_chars=5120,
+    placeholder="Type your sentence here..."
 )
 
-source_lang = st.selectbox(label="Source language", options=list(lang_id.keys()))
-target_lang = st.selectbox(label="Target language", options=list(lang_id.keys()))
+# Language selectors
+source_lang = st.selectbox("Select source language", sorted(lang_id.keys()))
+target_lang = st.selectbox("Select target language", sorted(lang_id.keys()))
 
-if st.button("Run"):
-    with st.spinner("Translating... please wait..."):
+# Translate Button
+if st.button("Translate"):
+    with st.spinner("Translating... Please wait"):
         time_start = time.time()
         tokenizer, model = load_model()
 
         src_lang = lang_id[source_lang]
         trg_lang = lang_id[target_lang]
+
         tokenizer.src_lang = src_lang
         with torch.no_grad():
             encoded_input = tokenizer(user_input, return_tensors="pt").to(device)
             generated_tokens = model.generate(
-                **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(trg_lang)
+                **encoded_input,
+                forced_bos_token_id=tokenizer.get_lang_id(trg_lang)
             )
             translated_text = tokenizer.batch_decode(
                 generated_tokens, skip_special_tokens=True
             )[0]
 
         time_end = time.time()
-        st.success(translated_text)
-        st.write(f"Computation time: {round((time_end - time_start), 3)} seconds")
-
+        st.success("Translation complete!")
+        st.markdown(f"**Translated Text:**\n\n{translated_text}")
+        st.caption(f"Time taken: {round(time_end - time_start, 2)} seconds")