Spaces:
Sleeping
Sleeping
File size: 2,040 Bytes
8056289 5ced387 6412832 8056289 6412832 5888fc0 203fe06 6412832 92c5432 6412832 5888fc0 e1fb94e 6412832 92c5432 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import streamlit as st
from transformers import pipeline, 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("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
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 summary..."):
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 summary and classification result
st.write("Summary:", summary)
st.write("Category:", predicted_label)
|