| import os |
| import tempfile |
| import fitz |
| import streamlit as st |
| import torch |
| from transformers import AutoConfig, AutoTokenizer, T5ForConditionalGeneration |
|
|
| |
| 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(): |
| |
| try: |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, legacy=False) |
| except Exception: |
| tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, legacy=False) |
|
|
| |
| 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] |
|
|
|
|
| |
| 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}") |
