Update app.py

app.py

@@ -1,7 +1,6 @@
-from transformers import pipeline
-from datasets import Dataset
 import streamlit as st
 import torch
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 
 # Set the background color and layout with set_page_config
 st.set_page_config(
@@ -10,7 +9,30 @@ st.set_page_config(
     layout="wide",
 )
 
-# Streamlit app setup
+# --- Model Loading with Caching ---
+# We use @st.cache_resource so the model loads only once, not on every user action.
+@st.cache_resource
+def load_model():
+    model_name = "Repleeka/mBART-tgj-final"
+    tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
+    model = MBartForConditionalGeneration.from_pretrained(model_name)
+    return tokenizer, model
+
+# Load model and tokenizer
+try:
+    tokenizer, model = load_model()
+    # Set the source language to English
+    tokenizer.src_lang = "en_XX" 
+except Exception as e:
+    st.error(f"Error loading model: {e}")
+
+# Determine device (GPU or CPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if 'model' in locals():
+    model = model.to(device)
+
+# --- Streamlit App Layout ---
+
 st.title(":repeat: English to Tagin Translator")
 st.markdown("Welcome to the English to Tagin Translator. :sparkles: Simply enter your text in English, and get the translation in Tagin instantly! :thumbsup:")
 
@@ -19,32 +41,31 @@ if 'text_input' not in st.session_state:
     st.session_state.text_input = ""
 text_input = st.text_area("Enter English text to translate", height=150, value=st.session_state.text_input)
 
-# Define your model from Hugging Face
-model_directory = "repleeka/eng-tagin-nmt"
-
-device = 0 if torch.cuda.is_available() else -1
-translation_pipeline = pipeline(
-    task="translation",
-    model="repleeka/eng-tagin-nmt",
-    tokenizer="repleeka/eng-tagin-nmt",
-    device=device
-)
-
 # Translate button
 if st.button("Translate", key="translate_button"):
     if text_input:
         with st.spinner("Translating... Please wait"):
-            # Prepare data for translation
-            sentences = [text_input]
-            data = Dataset.from_dict({"text": sentences})
-
-            # Apply translation
             try:
-                results = data.map(lambda x: {"translation": translation_pipeline(x["text"])})
-                result = results[0]["translation"][0]['translation_text']
+                # 1. Tokenize the input text
+                inputs = tokenizer(text_input, return_tensors="pt")
                 
-                # Capitalize the first letter of the result
-                result = result.capitalize()
+                # Move inputs to the correct device (GPU/CPU)
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+
+                # 2. Generate translation using the specific Tagin token ID
+                generated_tokens = model.generate(
+                    **inputs,
+                    forced_bos_token_id=tokenizer.convert_tokens_to_ids("<tgj_IN>"),
+                    num_beams=5,
+                    max_length=128,
+                )
+
+                # 3. Decode the generated tokens
+                result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+                
+                # Capitalize the first letter of the result for aesthetics
+                if result:
+                    result = result[0].upper() + result[1:]
 
                 # Display translation result with custom styling
                 st.markdown("#### Translated text:")
@@ -58,29 +79,23 @@ if st.button("Translate", key="translate_button"):
 # Clear input button
 if st.button("Clear Input"):
     st.session_state.text_input = ""
+    st.rerun() # Rerun to refresh the text area immediately
 
-st.markdown("""❗❗❗ **Please note:** The English-to-Tagin translator is still in its initial phase, so it may provide incorrect translations at times. Your understanding is appreciated!
+# --- Static Content ---
 
+st.markdown("""❗❗❗ **Please note:** The English-to-Tagin translator is still in its initial phase, so it may provide incorrect translations at times. Your understanding is appreciated!
 🤝 **For contributions or inquiries, feel free to contact me!**
-
 """)
 
-# Additional sections remain unchanged
 st.markdown("""### Tagin Language
-
 Tagin is a beautiful language spoken by the Tagin tribe and belongs to the Tani group of Sino-Tibetan languages. You'll mainly find Tagin speakers in the Upper Subansiri, Shiyomi, and in some parts Kara Dadi, and Kurung Kumey districts of Arunachal Pradesh, India. While about 63,000 (according to 2011 Census of India) people speak Tagin as their mother tongue, UNESCO has marked it as 'definitely endangered', which means it's at risk of disappearing. Unfortunately, very few written materials exist in Tagin, which makes it hard to study and preserve the language.
-
 As a small contribution to preserving this rich cultural heritage, I've developed this English-Tagin translator using the GinLish Corpus v0.1.0. By creating this digital tool, I hope to help keep the Tagin language alive and make it more accessible to both the Tagin community and language enthusiasts. This project is my way of giving back to society and helping protect an important piece of our cultural diversity.
-
 """)
 
 st.markdown("""### GinLish Corpus v0.1.0 (2024)
-
 I'm excited:satisfied: to share that I have created the GinLish Corpus v0.1.0, which is actually the first-ever collection of matched Tagin and English sentences. The corpus contains 60,000 carefully paired sentences that captures how these languages relate to each other. To build this, I translated English sentences from the Tatoeba website into Tagin and included traditional Tagin folk stories too.
 What makes this special is that I made sure to keep the true essence of the Tagin language alive in the translations. This means including Tagin sayings, cultural elements, and the unique way Tagin people express themselves. All this careful attention to detail makes the corpus really valuable for building translation tools, studying the language, and helping people learn Tagin. It's not just a simple word-for-word translation - it's a bridge between these two languages that respects and preserves Tagin's cultural identity.
-
 Good news:smiley: for researchers and language enthusiasts - I plan to release this dataset for non-commercial use once I complete my PhD!🎓 This way, others can also contribute to preserving and studying this beautiful language.
-
 """)
 
 # Sidebar for About and Contact information
@@ -88,13 +103,9 @@ st.sidebar.header("About the Developer")
 
 st.sidebar.markdown("""
 Hey there! 👋 
-
 I’m **Tungon Dugi**. 
-
 Right now, I’m doing my PhD in Computer Science and Engineering at NIT Arunachal Pradesh. 🎓 
-
 I’ve got a keen interest in Natural Language Processing (NLP), Machine Translation (MT), Deep Learning, and Linguistics. 
-
 💻✨ I love exploring how tech can help preserve and promote low-resource languages, especially my own language, **Tagin**! 🌍💬""")
 
 # Create some space between main sidebar content and footer
@@ -106,6 +117,6 @@ st.sidebar.caption("Contact: tungondugi@gmail.com")
 # Or using columns in sidebar:
 col1, col2 = st.sidebar.columns(2)
 with col1:
-    st.caption("© 2024")
+    st.caption("© 2026")
 with col2:
-    st.caption("v0.1.0")
+    st.caption("v0.1.1")