Update app.py

app.py

@@ -3,59 +3,63 @@ from transformers import pipeline
 import streamlit as st
 import fitz  # PyMuPDF for PDF text extraction
 
-# ------------------------------
-# Supported models
-MODEL_OPTIONS = {
-    "Long T5 (Scientific Simplifier)": "pszemraj/long-t5-tglobal-base-sci-simplify",
-    "T5 Base (General Simplifier)": "t5-base"
+# Streamlit UI setup
+st.set_page_config(page_title="Text Simplifier", layout="centered")
+st.title("🧠 Academic Text Simplifier")
+
+# Model selection
+model_options = {
+    "Mistral (Instruction-tuned)": "mistralai/Mistral-7B-Instruct-v0.1",
+    "T5 (Simplification finetuned)": "mrm8488/t5-base-finetuned-common_gen",
+    "BART (Paraphrasing/Simplification)": "tuner007/pegasus_paraphrase"
 }
 
-@st.cache_resource
-def load_model(model_name):
-    return pipeline("summarization", model=model_name)
+model_choice = st.selectbox("Choose a simplification model:", list(model_options.keys()))
+model_name = model_options[model_choice]
+
+@st.cache_resource(show_spinner=True)
+def load_model(name):
+    task = "text2text-generation" if "t5" in name.lower() or "pegasus" in name.lower() else "text-generation"
+    return pipeline(task, model=name)
+
+simplifier = load_model(model_name)
 
-def simplify_text(text, simplifier, model_name):
+def simplify_text(text):
     try:
-        # T5 expects a "summarize: " prefix
         if "t5" in model_name.lower():
-            text = "summarize: " + text
+            prompt = f"simplify: {text}"
+        elif "mistral" in model_name.lower() or "instruct" in model_name.lower():
+            prompt = f"Rewrite the following text using simpler vocabulary and structure:\n{text}"
+        elif "pegasus" in model_name.lower():
+            prompt = f"paraphrase: {text}"
+        else:
+            prompt = text
 
-        simplified = simplifier(text, max_length=256, min_length=30, do_sample=False)
-        return simplified[0]['summary_text']
+        output = simplifier(prompt, max_length=256, min_length=30, do_sample=False)[0]
+        return output.get('summary_text') or output.get('generated_text') or "(No output)"
     except Exception as e:
         return f"Error simplifying text: {e}"
 
-def extract_text_from_pdf(uploaded_file):
-    with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:
-        text = "\n".join(page.get_text() for page in doc)
+def extract_text_from_pdf(pdf_file):
+    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
+    text = "\n".join(page.get_text("text") for page in doc)
     return text
 
-# ------------------------------
-# Streamlit UI
-st.set_page_config(page_title="Jargon Simplifier", layout="centered")
-st.title("🧠 Jargon to Simple: Academic Text Simplifier")
-
-selected_model_name = st.selectbox("Choose a simplification model:", list(MODEL_OPTIONS.keys()))
-model_id = MODEL_OPTIONS[selected_model_name]
-simplifier = load_model(model_id)
-
+# Input options
 option = st.radio("Choose input type:", ("Text Input", "Upload PDF"))
 
 if option == "Text Input":
-    user_text = st.text_area("Enter complex academic text:")
-    if st.button("Simplify") and user_text.strip():
-        with st.spinner("Simplifying..."):
-            simplified_output = simplify_text(user_text, simplifier, model_id)
-        st.text_area("Simplified Output:", value=simplified_output, height=200)
+    user_text = st.text_area("Enter your complex academic text here:", height=200)
+    if st.button("Simplify Text") and user_text:
+        simplified_text = simplify_text(user_text)
+        st.text_area("🔽 Simplified Text:", simplified_text, height=200)
 
 elif option == "Upload PDF":
-    uploaded_file = st.file_uploader("Upload a PDF file:", type=["pdf"])
-    if uploaded_file:
-        with st.spinner("Extracting and simplifying text..."):
-            extracted_text = extract_text_from_pdf(uploaded_file)
-            truncated_text = extracted_text[:2000]  # Trim for model input
-            simplified_output = simplify_text(truncated_text, simplifier, model_id)
-        st.text_area("Simplified Output:", value=simplified_output, height=200)
+    uploaded_file = st.file_uploader("Upload a PDF document", type=["pdf"])
+    if uploaded_file and st.button("Simplify Extracted Text"):
+        extracted_text = extract_text_from_pdf(uploaded_file)
+        simplified_text = simplify_text(extracted_text[:2000])  # limit for performance
+        st.text_area("🔽 Simplified Text from PDF:", simplified_text, height=200)
 
 st.markdown("---")
 st.markdown("Made with ❤️ by Harshitha")