Emily666666's picture
Update app.py
44eee89 verified
raw
history blame
No virus
2.37 kB
import streamlit as st
from transformers import pipeline
# Load the text summarization pipeline
try:
summarizer = pipeline("summarization", model="syndi-models/titlewave-t5-base")
summarizer_loaded = True
except ValueError as e:
st.error(f"Error loading summarization model: {e}")
summarizer_loaded = False
# Load the news classification pipeline
model_name = "elozano/bert-base-cased-news-category"
try:
classifier = pipeline("text-classification", model=model_name, return_all_scores=True)
classifier_loaded = True
except ValueError as e:
st.error(f"Error loading classification model: {e}")
classifier_loaded = False
# Streamlit app title
st.title("Summarization and News Classification")
# Tab layout
tab1, tab2 = st.tabs(["Text Summarization", "News Classification"])
with tab1:
st.header("Text Summarization")
# Input text for summarization
text_to_summarize = st.text_area("Enter text to summarize:", "")
if st.button("Summarize"):
if summarizer_loaded and text_to_summarize:
try:
# Perform text summarization
summary = summarizer(text_to_summarize, max_length=130, min_length=30, do_sample=False)
# Display the summary result
st.write("Summary:", summary[0]['summary_text'])
except Exception as e:
st.error(f"Error during summarization: {e}")
else:
st.warning("Please enter text to summarize and ensure the model is loaded.")
with tab2:
st.header("News Classification")
# Input text for news classification
text_to_classify = st.text_area("Enter text to classify:", "")
if st.button("Classify"):
if classifier_loaded and text_to_classify:
try:
# Perform news classification
results = classifier(text_to_classify)[0]
# Find the category with the highest score
max_score = max(results, key=lambda x: x['score'])
st.write("Text:", text_to_classify)
st.write("Category:", max_score['label'])
st.write("Score:", max_score['score'])
except Exception as e:
st.error(f"Error during classification: {e}")
else:
st.warning("Please enter text to classify and ensure the model is loaded.")