Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification | |
import torch | |
# Define the summarization pipeline | |
summarizer_ntg = pipeline("summarization", model="mrm8488/t5-base-finetuned-summarize-news") | |
# Load the tokenizer and model for classification | |
tokenizer_bb = AutoTokenizer.from_pretrained("your-username/your-model-name") | |
model_bb = AutoModelForSequenceClassification.from_pretrained("your-username/your-model-name") | |
# 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'] | |
# 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) |