Update app.py

app.py

@@ -1,16 +1,23 @@
 import streamlit as st
 from transformers import pipeline
 
-# Load the summarization model
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+# Caching model loading to optimize memory usage
+@st.cache_resource
+def load_summarizer():
+    return pipeline("summarization", model="facebook/bart-large-cnn")
 
-# Load the grammar correction model
-grammar_correction_pipe = pipeline("text2text-generation", model="pszemraj/flan-t5-large-grammar-synthesis")
+@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)[0]['generated_text']
+        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."
@@ -18,7 +25,7 @@ def correct_grammar(user_input):
 # Function for text summarization
 def summarize_text(user_input):
     if user_input.strip():
-        summary = summarizer(user_input, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+        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."
@@ -38,7 +45,7 @@ task = st.selectbox("Choose a task", ["Summarize Text", "Correct Grammar"])
 # Input component for text
 user_input = st.text_area("Enter your text here:")
 
-# Process and display output based on selected task
+# Submit button
 if st.button("Submit"):
     if task == "Summarize Text":
         output = correct_and_summarize(user_input)  # Correct grammar, then summarize
@@ -46,4 +53,4 @@ if st.button("Submit"):
         output = correct_grammar(user_input)  # Only correct grammar
     
     # Display the output
-    st.text_area("Output", output, height=200, disabled=True)
+    st.text_area("Output", output, height=200)