Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,32 +2,58 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Model name
-MODEL_NAME = "AbdullahAlnemr1/flan-t5-summarizer"
+# ---- Configuration ----
+MODEL_NAME = "AbdullahAlnemr1/flan-t5-summarizer"  
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+# ---- Load model and tokenizer ----
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
+    model.to(device)
+    return tokenizer, model
 
-st.title("Text Summarizer(Encoder-Decoder)")
+tokenizer, model = load_model()
 
-input_text = st.text_area("Enter text to summarize:", height=200)
+# ---- Streamlit App ----
+st.title("Text Summarizer")
+st.write("Generate concise summaries from long articles using a fine-tuned Encoder–Decoder model.")
 
-max_new_tokens = st.slider("Max summary length", min_value=20, max_value=200, value=100)
+# ---- Input Area ----
+article = st.text_area("Enter the article or passage to summarize:", height=250)
 
+# ---- Parameters ----
+max_input_len = 512
+max_output_len = 150
+
+# ---- Generate Summary ----
 if st.button("Generate Summary"):
-    if input_text.strip() == "":
+    if not article.strip():
         st.warning("Please enter some text to summarize.")
     else:
-        # Tokenize input
-        inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
-        # Generate summary
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_new_tokens=max_new_tokens,
-            num_beams=4,
-            early_stopping=True
-        )
-        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.subheader("Summary:")
+        with st.spinner("Generating summary..."):
+            inputs = tokenizer(
+                article,
+                return_tensors="pt",
+                max_length=max_input_len,
+                truncation=True
+            ).to(device)
+
+            summary_ids = model.generate(
+                **inputs,
+                max_length=max_output_len,
+                num_beams=4,
+                length_penalty=2.0,
+                early_stopping=True
+            )
+
+            summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+        # ---- Output ----
+        st.subheader("Generated Summary:")
         st.write(summary)
+
+# ---- Footer ----
+st.markdown("---")
+st.markdown("Model powered by Transformers | Streamlit App by Ali Hamza")