File size: 3,361 Bytes
a92f235 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os
import tempfile
import fitz # PyMuPDF
import streamlit as st
import torch
from transformers import AutoConfig, AutoTokenizer, T5ForConditionalGeneration
# Page Configuration
st.set_page_config(
page_title="RTI & Civic PDF Simplifier",
page_icon="📄",
layout="centered",
)
st.title("📄 Civic & RTI PDF Simplifier")
st.markdown(
"Upload bureaucratic, legal, or RTI PDF documents to generate concise,"
" plain-language bullet points."
)
MODEL_REPO = "CaseLoop/pdf"
FALLBACK_BASE = "google/flan-t5-base"
@st.cache_resource(show_spinner="Loading model weights...")
def load_model_and_tokenizer():
# Tokenizer fallback
try:
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, legacy=False)
except Exception:
tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, legacy=False)
# Model loader with architecture fallback
try:
model = T5ForConditionalGeneration.from_pretrained(MODEL_REPO)
except Exception:
try:
config = AutoConfig.from_pretrained(FALLBACK_BASE)
model = T5ForConditionalGeneration.from_pretrained(
MODEL_REPO, config=config
)
except Exception:
model = T5ForConditionalGeneration.from_pretrained(FALLBACK_BASE)
model.eval()
return tokenizer, model
tokenizer, model = load_model_and_tokenizer()
def extract_chunks(
pdf_path: str, chunk_size: int = 180, max_chunks: int = 6
) -> list[str]:
"""Extracts lines from a PDF and chunks words for FLAN-T5 input limits."""
doc = fitz.open(pdf_path)
lines = []
for page in doc:
text = page.get_text("text")
lines.extend([
line.strip()
for line in text.split("\n")
if len(line.strip().split()) >= 5
])
words = " ".join(lines).split()
if not words:
return []
chunks = [
" ".join(words[i : i + chunk_size])
for i in range(0, len(words), chunk_size)
]
return chunks[:max_chunks]
# UI File Uploader
uploaded_file = st.file_uploader(
"Choose a PDF file", type=["pdf"], help="Upload an RTI reply or civic notice"
)
if uploaded_file is not None:
if st.button("Simplify Document", type="primary"):
with st.spinner("Parsing document and running simplification..."):
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(uploaded_file.read())
tmp_path = tmp.name
chunks = extract_chunks(tmp_path)
if os.path.exists(tmp_path):
os.remove(tmp_path)
if not chunks:
st.warning(
"No extractable text found in this PDF. It might be a scanned image."
)
else:
st.subheader("Key Takeaways in Plain Language:")
for idx, chunk in enumerate(chunks, 1):
prompt = f"simplify bureaucratic text into simple language: {chunk}"
inputs = tokenizer(
prompt, return_tensors="pt", max_length=512, truncation=True
)
with torch.no_grad():
output_tokens = model.generate(
**inputs,
max_length=120,
min_length=20,
num_beams=4,
no_repeat_ngram_size=3,
repetition_penalty=2.5,
early_stopping=True,
)
summary = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
st.markdown(f"**{idx}.** {summary}") |