from transformers import pipeline
from datasets import Dataset
import streamlit as st
import torch

# Set the background color and layout with set_page_config
st.set_page_config(
    page_title="English to Tagin Translator",
    page_icon=":repeat:",
    layout="wide",
)

# Streamlit app setup
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:")

# Text input
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']
                
                # Capitalize the first letter of the result
                result = result.capitalize()

                # Display translation result with custom styling
                st.markdown("#### Translated text:")
                st.markdown(f'<h2 class="result-text">{result}</h2>', unsafe_allow_html=True)
                
            except Exception as e:
                st.error(f"Translation error: {e}")
    else:
        st.warning("Please enter text to translate.")

# Clear input button
if st.button("Clear Input"):
    st.session_state.text_input = ""

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
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
st.sidebar.markdown("<br>" * 5, unsafe_allow_html=True)
st.sidebar.markdown("---")
st.sidebar.caption("Made with ❤️ by Tungon Dugi")
st.sidebar.caption("Contact: tungondugi@gmail.com")

# Or using columns in sidebar:
col1, col2 = st.sidebar.columns(2)
with col1:
    st.caption("© 2024")
with col2:
    st.caption("v0.1.0")