import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

def main():
    st.set_page_config(page_title="LinguaSphere: Aya by CohereAI", layout="wide")

    st.title("LinguaSphere: Aya by CohereAI")

    # Load translation model and tokenizer
    translation_checkpoint = "CohereForAI/aya-101"
    translation_tokenizer = AutoTokenizer.from_pretrained(translation_checkpoint)
    translation_model = AutoModelForSeq2SeqLM.from_pretrained(translation_checkpoint)

    # Load sentiment analysis model
    sentiment_analysis_pipeline = pipeline("sentiment-analysis", model="CohereForAI/aya-101")

    # Load named entity recognition (NER) model
    ner_pipeline = pipeline("ner", model="CohereForAI/aya-101")

    # Load summarization model
    summarization_pipeline = pipeline("summarization", model="CohereForAI/aya-101")

    # Sidebar options
    st.sidebar.title("Options")
    show_model_info = st.sidebar.checkbox("Show Model Information")
    show_task_description = st.sidebar.checkbox("Show Task Description")
    show_about = st.sidebar.checkbox("About")

    if show_about:
        st.sidebar.subheader("About")
        st.sidebar.markdown(
            "LinguaSphere is a Streamlit app powered by Aya, a multilingual model developed by CohereAI."
        )

    if show_model_info:
        st.sidebar.subheader("Model Information")
        st.sidebar.markdown(
            """
            The Aya model is a massively multilingual generative language model developed by Cohere For AI. 
            It is capable of performing various natural language processing tasks such as translation, sentiment analysis, 
            named entity recognition, and summarization.
            """
        )

    if show_task_description:
        st.sidebar.subheader("Task Descriptions")
        st.sidebar.markdown(
            """
            - **Translation:** Translate text from one language to another.
            - **Text Generation:** Generate text based on a prompt.
            - **Sentiment Analysis:** Analyze the sentiment of text.
            - **Named Entity Recognition:** Identify named entities in text.
            - **Summarization:** Generate a summary of text.
            """
        )

    task = st.selectbox(
        "Select Task",
        [
            "Translation",
            "Text Generation",
            "Sentiment Analysis",
            "Named Entity Recognition",
            "Summarization",
        ],
    )

    if task == "Translation":
        source_language = st.selectbox(
            "Source Language", ["English", "Turkish", "Hindi"]
        )
        target_language = st.selectbox(
            "Target Language", ["Nepali", "Turkish", "English"]
        )
        source_text = st.text_area("Enter text in " + source_language, "")

        if st.button("Translate"):
            source_text = source_text.strip()
            if source_language == "English":
                source_text = "Translate to " + target_language + ": " + source_text
            else:
                source_text = source_text + "।"
            inputs = translation_tokenizer.encode(source_text, return_tensors="pt")
            outputs = translation_model.generate(inputs, max_length=128)
            translated_text = translation_tokenizer.decode(
                outputs[0], skip_special_tokens=True
            )
            st.write("Translated text in " + target_language + ":", translated_text)

    elif task == "Text Generation":
        prompt = st.text_area("Enter prompt for text generation", "")
        language = st.selectbox("Select Language", ["English", "Turkish", "Hindi"])
        if st.button("Generate"):
            prompt = prompt.strip()
            if language == "English":
                prompt = "Generate text: " + prompt
            elif language == "Hindi":
                prompt = prompt + "।"
            inputs = translation_tokenizer.encode(prompt, return_tensors="pt")
            outputs = translation_model.generate(inputs, max_length=128)
            generated_text = translation_tokenizer.decode(
                outputs[0], skip_special_tokens=True
            )
            st.write("Generated text:", generated_text)

    elif task == "Sentiment Analysis":
        text = st.text_area("Enter text for sentiment analysis", "")
        if st.button("Analyze Sentiment"):
            result = sentiment_analysis_pipeline(text)
            st.write("Sentiment:", result[0]["label"])

    elif task == "Named Entity Recognition":
        text = st.text_area("Enter text for named entity recognition", "")
        if st.button("Recognize Entities"):
            entities = ner_pipeline(text)
            st.write("Entities:")
            for entity in entities:
                st.write(entity)

    elif task == "Summarization":
        text = st.text_area("Enter text for summarization", "")
        if st.button("Summarize"):
            summary = summarization_pipeline(text)
            st.write("Summary:", summary[0]["summary_text"])


if __name__ == "__main__":
    main()