import streamlit as st
from transformers import pipeline

# Load the text summarization pipeline
try:
    model3_p1 = pipeline("summarization", model="syndi-models/titlewave-t5-base")
except ValueError as e:
    st.error(f"Error loading summarization model: {e}")

# Load the news classification pipeline
model_name2_p2 = "elozano/bert-base-cased-news-category"
try:
    classifier = pipeline("text-classification", model=model_name2_p2, return_all_scores=True)
except ValueError as e:
    st.error(f"Error loading classification model: {e}")

# Streamlit app title
st.title("Summarization and News Classification")

# Tab layout
tab1, tab2 = st.tabs(["Text Summarization", "News Classification"])

with tab1:
    st.header("Text Summarization")
    # Input text for summarization
    text_to_summarize = st.text_area("Enter text to summarize:", "")
    if st.button("Summarize") and 'model3_p1' in globals():
        # Perform text summarization
        summary = model3_p1(text_to_summarize, max_length=130, min_length=30, do_sample=False)
        # Display the summary result
        st.write("Summary:", summary[0]['summary_text'])

with tab2:
    st.header("News Classification")
    # Input text for news classification
    text_to_classify = st.text_area("Enter text to classify:", "")
    if st.button("Classify") and 'classifier' in globals():
        # Perform news classification
        results = classifier(text_to_classify)[0]
        # Display the classification result
        max_score = float('-inf')
        max_label = ''
        for result in results:
            if result['score'] > max_score:
                max_score = result['score']
                max_label = result['label']
        st.write("Text:", text_to_classify)
        st.write("Category:", max_label)
        st.write("Score:", max_score)