import streamlit as st
from transformers import pipeline

st.set_page_config(
    page_title="Multilingual Poem Generator",
    page_icon="🎭",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.title("🌍 Multilingual Poem Generator")
st.write("Generate poems in English and 7 Indian regional languages.")

# Language options
languages = {
    "English": "gpt2",  # Default GPT-2 (English)
    "Hindi": "ai4bharat/IndicGPT-hindi",
    "Bengali": "ai4bharat/IndicGPT-bengali",
    "Tamil": "ai4bharat/IndicGPT-tamil",
    "Telugu": "ai4bharat/IndicGPT-telugu",
    "Marathi": "ai4bharat/IndicGPT-marathi",
    "Gujarati": "ai4bharat/IndicGPT-gujarati",
    "Kannada": "ai4bharat/IndicGPT-kannada"
}

selected_lang = st.selectbox("Choose a language:", list(languages.keys()), index=0)

prompt_input = st.text_area("Enter a few words or sentences to start:",
                            value="Once upon a time,",
                            height=100)

col1, empty_col, col2 = st.columns([0.7, 0.3, 1.0])
with col1:
    max_length = st.slider("Max Output Tokens:", 10, 200, 75)
    temperature = st.slider("Temperature:", 0.1, 1.5, 0.8)
    top_k = st.slider("Top-k:", 1, 100, 50)
    top_p = st.slider("Top-p:", 0.1, 1.0, 0.95)
    repetition_penalty = st.slider("Repetition Penalty:", 1.0, 2.0, 1.2)

with empty_col:
    st.empty()

with col2:
    if st.button("Generate Poem"):
        with st.spinner(f"Generating poem in {selected_lang}..."):
            try:
                # Load model for selected language
                generator = pipeline(
                    "text-generation",
                    model=languages[selected_lang]
                )

                # Generate poem
                results = generator(
                    prompt_input,
                    max_new_tokens=max_length,
                    do_sample=True,
                    top_k=top_k,
                    top_p=top_p,
                    temperature=temperature,
                    repetition_penalty=repetition_penalty
                )

                generated_text = results[0]['generated_text']
                st.subheader(f"Generated Poem in {selected_lang}:")
                for line in generated_text.split('\n'):
                    st.write(line)

            except Exception as e:
                st.error(f"An error occurred: {str(e)}")