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

# Load the tokenizer and model for classification
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")

tokenizer_bb = AutoTokenizer.from_pretrained("Lauraayu/News_Classification_Model")
model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classification_Model")

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

# Text input for user to enter the news article text
article = st.text_area("News Article", height=300)

def summarize(text, max_length=150):
    input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
    generated_ids = model.generate(input_ids=input_ids, num_beams=2, max_length=max_length, repetition_penalty=2.5, length_penalty=1.0, early_stopping=True)
    preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
    return preds[0]

    
# Perform summarization and classification when the user clicks the "Classify" button
if st.button("Classify"):
    # Perform text summarization
    with st.spinner("Generating category..."):
            summary = summarize(article)
    
    # Tokenize the summarized text
    inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
    
    # 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 classification result
    st.write("Summary:", summary)
    st.write("Category:", predicted_label)