import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch

# Define the summarization pipeline
try:
    summarizer_ntg = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-summarize-news")
except Exception as e:
    st.error(f"Error loading summarization model: {e}")

# Load the tokenizer and model for classification
try:
    tokenizer_bb = AutoTokenizer.from_pretrained("Lauraayu/News_Classi_Model")
    model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classi_Model")
except Exception as e:
    st.error(f"Error loading classification model or tokenizer: {e}")

# Streamlit application title
st.title("News Article Summarizer and Classifier")
st.write("Enter a news article text to get its summary and category.")

# Text input for user to enter the news article text
text = st.text_area("Enter the news article text here:")

# Perform summarization and classification when the user clicks the "Classify" button
if st.button("Classify"):
    if not text:
        st.error("Please enter some text to classify.")
    else:
        try:
            # Perform text summarization
            summary = summarizer_ntg(text)[0]['summary_text']
            
            # Tokenize the summarized text
            inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
            
            # Move inputs and model to the same device (GPU or CPU)
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            inputs = {k: v.to(device) for k, v in inputs.items()}
            model_bb.to(device)
            
            # Perform text classification
            with torch.no_grad():
                outputs = model_bb(**inputs)
            
            # Get the predicted label
            predicted_label_id = torch.argmax(outputs.logits, dim=-1).item()
            label_mapping = model_bb.config.id2label
            predicted_label = label_mapping[predicted_label_id]
            
            # Display the summary and classification result
            st.write("Summary:", summary)
            st.write("Category:", predicted_label)
        except Exception as e:
            st.error(f"Error during summarization or classification: {e}")