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 Tawra Translator",
    page_icon=":repeat:",
    layout="wide",
)

# Streamlit app setup
st.title(":repeat: English to Tawra Translator")
st.markdown("Welcome to the English to Tawra Translator. :sparkles: Simply enter your text in English, and get the translation in Tawra 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-taw-nmt"

device = 0 if torch.cuda.is_available() else -1
translation_pipeline = pipeline(
    task="translation",
    model="repleeka/eng-taw-nmt",
    tokenizer="repleeka/eng-taw-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}</2>', unsafe_allow_html=True)
                # st.markdown(result)
                
            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 = ""