import streamlit as st
from transformers import pipeline

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

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

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

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

with tab1:
    st.header("Question Summarization")
    # Input text for summarization
    text_to_summarize = st.text_area("Enter long question to summarize:", "")
    if st.button("Summarize"):
        if summarizer_loaded and text_to_summarize:
            try:
                # Perform text summarization
                summary = summarizer(text_to_summarize, max_length=130, min_length=30, do_sample=False)
                # Display the summary result
                st.write("Summary:", summary[0]['summary_text'])
            except Exception as e:
                st.error(f"Error during summarization: {e}")
        else:
            st.warning("Please enter text to summarize and ensure the model is loaded.")

with tab2:
    st.header("Questions Classification")
    # Input text for Question classification
    text_to_classify = st.text_area("Enter question to classify:", "")
    if st.button("Classify"):
        if classifier_loaded and text_to_classify:
            try:
                # Perform uestion classification
                results = classifier(text_to_classify)[0]
                # Find the category with the highest score
                max_score = max(results, key=lambda x: x['score'])
                st.write("Text:", text_to_classify)
                st.write("Category:", max_score['label'])
                st.write("Score:", max_score['score'])
            except Exception as e:
                st.error(f"Error during classification: {e}")
        else:
            st.warning("Please enter text to classify and ensure the model is loaded.")