import streamlit as st
from transformers import pipeline

# Caching model loading to optimize memory usage
@st.cache_resource
def load_summarizer():
    return pipeline("summarization", model="facebook/bart-large-cnn")

@st.cache_resource
def load_grammar_correction_pipe():
    return pipeline("text2text-generation", model="pszemraj/flan-t5-large-grammar-synthesis")

# Initialize models using the cached loading functions
summarizer = load_summarizer()
grammar_correction_pipe = load_grammar_correction_pipe()

# Function for grammar correction
def correct_grammar(user_input):
    if user_input.strip():
        corrected_text = grammar_correction_pipe(user_input, max_length=256)[0]['generated_text']
        return corrected_text
    else:
        return "Please enter some text for grammar correction."

# Function for text summarization
def summarize_text(user_input):
    if user_input.strip():
        summary = summarizer(user_input, max_length=100, min_length=30, do_sample=False)[0]['summary_text']
        return summary
    else:
        return "Please enter some text to summarize."

# Function to combine grammar correction and summarization
def correct_and_summarize(user_input):
    corrected_text = correct_grammar(user_input)  # First correct the grammar
    summary = summarize_text(corrected_text)  # Then summarize the corrected text
    return summary

# Streamlit UI setup
st.title("Text Summarization and Grammar Correction Assistant")

# Dropdown to select task
task = st.selectbox("Choose a task", ["Summarize Text", "Correct Grammar"])

# Input component for text
user_input = st.text_area("Enter your text here:")

# Submit button
if st.button("Submit"):
    if task == "Summarize Text":
        output = correct_and_summarize(user_input)  # Correct grammar, then summarize
    elif task == "Correct Grammar":
        output = correct_grammar(user_input)  # Only correct grammar
    
    # Display the output
    st.text_area("Output", output, height=200)