Summarization / app.py
llmahmad's picture
Update app.py
97814f0 verified
raw
history blame
1.54 kB
import os
import streamlit as st
from transformers import PegasusTokenizer, PegasusForConditionalGeneration
# Install necessary packages (only needed if not already installed)
os.system('pip install streamlit transformers torch')
# Load the model and tokenizer
model_path = '.' # Path to the current directory where files are located
tokenizer = PegasusTokenizer.from_pretrained(model_path)
model = PegasusForConditionalGeneration.from_pretrained(model_path)
def summarize_text(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="longest")
summary_ids = model.generate(inputs["input_ids"], max_length=150, min_length=30, do_sample=False)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
st.title("Text Summarization with Fine-Tuned Model")
st.write("Enter text to generate a summary using the fine-tuned summarization model.")
text = st.text_area("Input Text", height=200)
if st.button("Summarize"):
if text:
with st.spinner("Summarizing..."):
summary = summarize_text(text)
st.success("Summary Generated")
st.write(summary)
else:
st.warning("Please enter some text to summarize.")
if __name__ == "__main__":
st.set_option('deprecation.showfileUploaderEncoding', False)
st.markdown(
"""
<style>
.reportview-container {
flex-direction: row;
justify-content: center;
}
</style>
""",
unsafe_allow_html=True
)