import streamlit as st
from transformers import pipeline

# Set up the app title
st.title("Language Learning App with Hugging Face")

# Add a sidebar for user input
st.sidebar.header("Settings")
task = st.sidebar.selectbox(
    "Choose a task",
    ["Translation", "Text Generation", "Sentiment Analysis"]
)

# Load the Hugging Face pipeline based on the selected task
if task == "Translation":
    model_name = st.sidebar.selectbox(
        "Choose a translation model",
        ["Helsinki-NLP/opus-mt-en-fr", "Helsinki-NLP/opus-mt-fr-en"]
    )
    st.sidebar.write(f"Selected model: {model_name}")
    try:
        translator = pipeline("translation", model=model_name)
    except Exception as e:
        st.error(f"Error loading translation model: {e}")

elif task == "Text Generation":
    model_name = st.sidebar.selectbox(
        "Choose a text generation model",
        ["gpt2", "EleutherAI/gpt-neo-125M"]
    )
    st.sidebar.write(f"Selected model: {model_name}")
    try:
        generator = pipeline("text-generation", model=model_name)
    except Exception as e:
        st.error(f"Error loading text generation model: {e}")

elif task == "Sentiment Analysis":
    model_name = st.sidebar.selectbox(
        "Choose a sentiment analysis model",
        ["distilbert-base-uncased-finetuned-sst-2-english"]
    )
    st.sidebar.write(f"Selected model: {model_name}")
    try:
        analyzer = pipeline("sentiment-analysis", model=model_name)
    except Exception as e:
        st.error(f"Error loading sentiment analysis model: {e}")

# Main app functionality
st.header(f"{task} Task")

if task == "Translation":
    text = st.text_area("Enter text to translate", "Hello, how are you?")
    if st.button("Translate"):
        try:
            translation = translator(text)
            st.write("Translation:")
            st.write(translation[0]['translation_text'])
        except Exception as e:
            st.error(f"Error during translation: {e}")

elif task == "Text Generation":
    prompt = st.text_area("Enter a prompt", "Once upon a time")
    max_length = st.slider("Max length", 10, 100, 50)
    if st.button("Generate Text"):
        try:
            generated_text = generator(prompt, max_length=max_length)
            st.write("Generated Text:")
            st.write(generated_text[0]['generated_text'])
        except Exception as e:
            st.error(f"Error during text generation: {e}")

elif task == "Sentiment Analysis":
    text = st.text_area("Enter text for sentiment analysis", "I love learning new languages!")
    if st.button("Analyze Sentiment"):
        try:
            sentiment = analyzer(text)
            st.write("Sentiment Analysis Result:")
            st.write(sentiment[0])
        except Exception as e:
            st.error(f"Error during sentiment analysis: {e}")