Lauraayu's picture
Update app.py
2f910ac verified
raw
history blame
1.41 kB
import streamlit as st
from transformers import pipeline
import torch
# Define pipelines
summarizer_ntg = pipeline("summarization", model="mrm8488/t5-base-finetuned-summarize-news")
model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classi_Model")
# 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"):
# Perform text summarization
summary = summarizer_ntg(text)[0]['summary_text']
# Move inputs and model to the same device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
summary = {k: v.to(device) for k, v in inputs.items()}
model_bb.to(device)
# Perform text classification
with torch.no_grad():
outputs = model_bb(**summary)
# 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)