Update app.py

app.py

@@ -1,37 +1,39 @@
 import gradio as gr
-import PyPDF2
-import docx2txt
 from transformers import pipeline
+import docx2txt
+import pdfplumber
 
-# Load summarization pipeline
-summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+# Load summarization model
+summarizer = pipeline("summarization", model="facebook/bart-base")
 
 def extract_text_from_file(file):
     if file.name.endswith(".pdf"):
-        reader = PyPDF2.PdfReader(file)
-        return " ".join([page.extract_text() for page in reader.pages if page.extract_text()])
+        with pdfplumber.open(file.name) as pdf:
+            return "\n".join([page.extract_text() for page in pdf if page.extract_text()])
     elif file.name.endswith(".docx"):
         return docx2txt.process(file.name)
-    else:
-        return "Unsupported file type."
+    return ""
 
-def generate_summary(file):
-    try:
-        text = extract_text_from_file(file)
-        if not text:
-            return "No text found in file."
-        chunks = [text[i:i+1024] for i in range(0, len(text), 1024)]
-        summaries = [summarizer(chunk)[0]['summary_text'] for chunk in chunks]
-        return "\n\n".join(summaries)
-    except Exception as e:
-        return f"Error processing file: {str(e)}"
+def generate_summary(files):
+    summaries = []
+    for file in files:
+        raw_text = extract_text_from_file(file)
+        if raw_text:
+            chunks = [raw_text[i:i+1000] for i in range(0, len(raw_text), 1000)]
+            summary = ""
+            for chunk in chunks:
+                result = summarizer(chunk, max_length=130, min_length=30, do_sample=False)
+                summary += result[0]['summary_text'] + "\n"
+            summaries.append(f"### Summary for {file.name}:\n{summary}")
+        else:
+            summaries.append(f"Could not extract text from {file.name}")
+    return "\n\n".join(summaries)
 
-iface = gr.Interface(
-    fn=generate_summary,
-    gr.File(label="Upload EHCPs", file_types=[".pdf", ".docx"], file_count="multiple")
-    outputs="text",
-    title="EHCP Summary Generator",
-    description="Upload an EHCP in PDF or Word format to get a structured summary using a transformer model."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## EHCP Document Summarizer\nUpload multiple EHCPs (PDF/DOCX) to generate summaries.")
+    file_input = gr.File(label="Upload EHCPs", file_types=[".pdf", ".docx"], file_count="multiple")
+    output = gr.Textbox(label="Summarised Output", lines=20)
+    submit_btn = gr.Button("Generate Summary")
+    submit_btn.click(fn=generate_summary, inputs=file_input, outputs=output)
 
-iface.launch()
+demo.launch()