File size: 1,987 Bytes
a5a1104 54036f6 a5a1104 54036f6 a5a1104 54036f6 a5a1104 54036f6 a5a1104 54036f6 a5a1104 9ff4d36 54036f6 |
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 49 50 51 52 53 54 55 56 57 |
import streamlit as st
from transformers import pipeline
# Caching model loading to optimize memory usage
@st.cache_resource
def load_summarizer():
return pipeline("summarization", model="facebook/bart-large-cnn")
@st.cache_resource
def load_grammar_correction_pipe():
return pipeline("text2text-generation", model="pszemraj/flan-t5-large-grammar-synthesis")
# Initialize models using the cached loading functions
summarizer = load_summarizer()
grammar_correction_pipe = load_grammar_correction_pipe()
# Function for grammar correction
def correct_grammar(user_input):
if user_input.strip():
corrected_text = grammar_correction_pipe(user_input, max_length=256)[0]['generated_text']
return corrected_text
else:
return "Please enter some text for grammar correction."
# Function for text summarization
def summarize_text(user_input):
if user_input.strip():
summary = summarizer(user_input, max_length=100, min_length=30, do_sample=False)[0]['summary_text']
return summary
else:
return "Please enter some text to summarize."
# Function to combine grammar correction and summarization
def correct_and_summarize(user_input):
corrected_text = correct_grammar(user_input) # First correct the grammar
summary = summarize_text(corrected_text) # Then summarize the corrected text
return summary
# Streamlit UI setup
st.title("Text Summarization and Grammar Correction Assistant")
# Dropdown to select task
task = st.selectbox("Choose a task", ["Summarize Text", "Correct Grammar"])
# Input component for text
user_input = st.text_area("Enter your text here:")
# Submit button
if st.button("Submit"):
if task == "Summarize Text":
output = correct_and_summarize(user_input) # Correct grammar, then summarize
elif task == "Correct Grammar":
output = correct_grammar(user_input) # Only correct grammar
# Display the output
st.text_area("Output", output, height=200)
|