Update app.py

app.py

@@ -1,21 +1,18 @@
 import streamlit as st
-import os
-import time
 import torch
 import logging
+import time
 from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
 
-# Set Streamlit page configuration
-st.set_page_config(page_title="M2M100 Translator")
+# Configure page
+st.set_page_config(page_title="🌐 Translator", page_icon="🌐")
 
-# 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.")
+# Device detection
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if device.type == "cpu":
+    logging.warning("⚠️ GPU not found — using CPU (translation may be slow).")
 
-# Language code mapping
+# Language mapping
 lang_id = {
     "Afrikaans": "af", "Amharic": "am", "Arabic": "ar", "Asturian": "ast",
     "Azerbaijani": "az", "Bashkir": "ba", "Belarusian": "be", "Bulgarian": "bg",
@@ -44,58 +41,58 @@ lang_id = {
     "Yiddish": "yi", "Yoruba": "yo", "Chinese": "zh", "Zulu": "zu",
 }
 
-# Cache the model and tokenizer using new API
+# Cache model/tokenizer loading
 @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
-    ).to(device)
+def load_model():
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")
+    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B").to(device)
     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)
-""")
+# Title
+st.title("🌍 M2M100 Language Translator")
+st.markdown("🔁 Translate text between **100+ languages** using Facebook's `M2M100` multilingual model.")
 
-# Input Text Area
+# Text input
 user_input = st.text_area(
-    "Enter text to translate:",
+    "✏️ Enter your text below:",
     height=200,
     max_chars=5120,
-    placeholder="Type your sentence here..."
+    placeholder="E.g. Hello, how are you?"
 )
 
-# Language selectors
-source_lang = st.selectbox("Select source language", sorted(lang_id.keys()))
-target_lang = st.selectbox("Select target language", sorted(lang_id.keys()))
+# Language selections (default: English → Hindi)
+col1, col2 = st.columns(2)
+with col1:
+    source_lang = st.selectbox("🌐 Source Language", sorted(lang_id.keys()), index=list(lang_id.keys()).index("English"))
+with col2:
+    target_lang = st.selectbox("🔁 Target Language", sorted(lang_id.keys()), index=list(lang_id.keys()).index("Hindi"))
 
 # Translate Button
-if st.button("Translate"):
+if st.button("🚀 Translate", disabled=(not user_input.strip())):
     with st.spinner("Translating... Please wait"):
-        time_start = time.time()
+        start = time.time()
         tokenizer, model = load_model()
 
-        src_lang = lang_id[source_lang]
-        trg_lang = lang_id[target_lang]
+        src = lang_id[source_lang]
+        tgt = lang_id[target_lang]
 
-        tokenizer.src_lang = src_lang
+        tokenizer.src_lang = src
         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 = tokenizer(user_input, return_tensors="pt").to(device)
+            output = model.generate(
+                **encoded,
+                forced_bos_token_id=tokenizer.get_lang_id(tgt)
             )
-            translated_text = tokenizer.batch_decode(
-                generated_tokens, skip_special_tokens=True
-            )[0]
+            result = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+
+        end = time.time()
+        st.success("✅ Translation complete!")
+        st.markdown("### 📝 Translated Text")
+        st.text_area("Output", value=result, height=150, disabled=True)
+        st.caption(f"⏱️ Time taken: {round(end - start, 2)} seconds")
 
-        time_end = time.time()
-        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")
+# Optional reset
+st.markdown("---")
+if st.button("🔄 Reset"):
+    st.experimental_rerun()