import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load the model and tokenizer
model_name = 'utrobinmv/t5_summary_en_ru_zh_base_2048'
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)

def summarize_text(text, prefix):
    src_text = prefix + text
    input_ids = tokenizer(src_text, return_tensors="pt")
    generated_tokens = model.generate(**input_ids)
    result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    return result[0]

st.title('Text Summarization App')

input_text = st.text_area("Enter the text to summarize:", height=300)

if st.button("Generate Summaries"):
    if input_text:
        title1 = summarize_text(input_text, 'summary: ')
        title2 = summarize_text(input_text, 'summary brief: ')
        st.write("### Title 1")
        st.write(title1)
        st.write("### Title 2")
        st.write(title2)
    else:
        st.warning("Please enter some text to summarize.")