File size: 3,357 Bytes
65284f8 191645f 65284f8 191645f | 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 | import streamlit as st
import os
import tempfile
from uuid import uuid4
from concurrent.futures import ThreadPoolExecutor, Future
from ingestion import ingest_pdf
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv(), override=False)
st.set_page_config(page_title="PDF Ingestion", layout="wide")
if "session_id" not in st.session_state:
st.session_state.session_id = str(uuid4())
if "ingestion_future" not in st.session_state:
st.session_state.ingestion_future = None
if "last_output" not in st.session_state:
st.session_state.last_output = None
st.title("Ingestion Layer")
st.caption("Upload a PDF to extract, normalize, chunk, and persist for later embedding.")
with st.sidebar:
st.header("Settings")
parser = st.selectbox("Parser", ["pypdf", "pdfplumber"], index=0)
chunk_size = st.number_input("Chunk size (chars)", min_value=500, max_value=10000, value=4000, step=100)
overlap = st.number_input("Overlap (chars)", min_value=0, max_value=2000, value=400, step=50)
uploaded = st.file_uploader("Upload PDF", type=["pdf"])
def _run_ingestion(tmp_path: str, sid: str, parser: str, chunk_size: int, overlap: int):
return ingest_pdf(tmp_path, sid, parser=parser, chunk_size=chunk_size, overlap=overlap)
col1, col2 = st.columns([2, 1])
with col1:
if uploaded is not None and st.button("Start Ingestion", type="primary"):
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(uploaded.read())
tmp_path = tmp.name
executor = ThreadPoolExecutor(max_workers=1)
future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, parser, int(chunk_size), int(overlap))
st.session_state.ingestion_future = future
st.info("Ingestion started. Processing in background...")
future: Future = st.session_state.ingestion_future
if future is not None:
if future.done():
try:
chunks, out_path = future.result()
st.session_state.last_output = (len(chunks), out_path)
st.success(f"Ingestion complete: {len(chunks)} chunks written to {out_path}")
if chunks:
st.subheader("Preview (first chunk)")
st.code(chunks[0].content[:2000])
except Exception as e:
st.error(f"Ingestion failed: {e}")
finally:
st.session_state.ingestion_future = None
else:
st.status("Processing...", state="running")
with col2:
st.subheader("Session")
st.code(st.session_state.session_id)
if st.session_state.last_output:
cnt, path = st.session_state.last_output
st.metric("Chunks", cnt)
st.write("Output:")
st.code(path)
if uploaded is not None and st.button("Retry with fallback parser"):
fb = "pdfplumber" if parser == "pypdf" else "pypdf"
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(uploaded.read())
tmp_path = tmp.name
executor = ThreadPoolExecutor(max_workers=1)
future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, fb, int(chunk_size), int(overlap))
st.session_state.ingestion_future = future
st.info(f"Retrying with {fb}...")
|