app.py

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

# 1. Page Config must be the first Streamlit command
st.set_page_config(
    page_title="English to Tawra Translator",
    page_icon=":repeat:",
    layout="wide",
)

# Custom CSS for the result text and general styling
st.markdown("""
    <style>
    .result-text {
        color: #1E88E5;
        background-color: #f0f2f6;
        padding: 20px;
        border-radius: 10px;
        border-left: 5px solid #1E88E5;
    }
    </style>
""", unsafe_allow_html=True)

# 2. Optimization: Cache the model loading to prevent reloading on every interaction
@st.cache_resource
def load_translator():
    model_id = "repleeka/eng-taw-nmt"
    # Determine device: use GPU if available, else CPU
    device = 0 if torch.cuda.is_available() else -1
    return pipeline(
        task="translation",
        model=model_id,
        tokenizer=model_id,
        device=device
    )

translation_pipeline = load_translator()

# 3. App UI 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:")

# 4. State Management for Input Clearing
if 'text_input' not in st.session_state:
    st.session_state.text_input = ""

def clear_text():
    st.session_state.text_input = ""

# Text area tied to session state
text_input = st.text_area(
    "Enter English text to translate", 
    height=150, 
    value=st.session_state.text_input,
    key="current_input" # Adding a key helps Streamlit track this specific widget
)

# Update session state with current input so it persists until "Clear" is clicked
st.session_state.text_input = text_input

col1, col2 = st.columns([1, 10])

with col1:
    translate_clicked = st.button("Translate", type="primary")

with col2:
    # Use a callback or direct state update for clearing
    if st.button("Clear Input"):
        clear_text()
        st.rerun() # Force rerun to clear the text area immediately

# 5. Translation Logic
if translate_clicked:
    if text_input.strip():
        with st.spinner("Translating... Please wait"):
            try:
                # Prepare data using the Datasets library
                sentences = [text_input]
                data = Dataset.from_dict({"text": sentences})

                # Apply translation
                # Fix: Hugging Face translation pipelines usually return a list of dicts directly
                # We use .map for batch processing consistency if you have many sentences
                results = data.map(lambda x: {"translation": translation_pipeline(x["text"])})
                
                # Extract the translation text
                # Accessing: Row 0 -> 'translation' column -> index 0 of results -> 'translation_text' key
                raw_result = results[0]["translation"][0]['translation_text']
                
                # Formatting
                final_result = raw_result.strip().capitalize()

                # Display result
                st.markdown("---")
                st.markdown("#### Translated text:")
                # Fixed the closing tag error in your HTML string (</h2> instead of </2>)
                st.markdown(f'<div class="result-text"><h2>{final_result}</h2></div>', unsafe_allow_html=True)
                
            except Exception as e:
                st.error(f"Translation error: {e}")
                # Optional: log error details for debugging
                # st.exception(e)
    else:
        st.warning("Please enter some text to translate.")