File size: 3,017 Bytes
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import json
import os
import re
from dataclasses import dataclass, asdict
from typing import List, Optional, Tuple
from pypdf import PdfReader
import pdfplumber
@dataclass
class Chunk:
content: str
source: str
page: Optional[int] = None
chapter: Optional[str] = None
def normalize_text(text: str) -> str:
# Remove control characters, normalize whitespace
text = text.encode("utf-8", errors="ignore").decode("utf-8", errors="ignore")
text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F]", " ", text)
text = re.sub(r"\s+", " ", text).strip()
return text
def extract_text_pypdf(file_path: str) -> List[Tuple[int, str]]:
pages: List[Tuple[int, str]] = []
reader = PdfReader(file_path)
for i, page in enumerate(reader.pages, start=1):
text = page.extract_text() or ""
pages.append((i, text))
return pages
def extract_text_pdfplumber(file_path: str) -> List[Tuple[int, str]]:
pages: List[Tuple[int, str]] = []
with pdfplumber.open(file_path) as pdf:
for i, page in enumerate(pdf.pages, start=1):
text = page.extract_text() or ""
pages.append((i, text))
return pages
def chunk_text(text: str, chunk_size: int = 4000, overlap: int = 400) -> List[str]:
# Simple recursive-like character splitter
if chunk_size <= 0:
return [text]
chunks: List[str] = []
start = 0
n = len(text)
while start < n:
end = min(start + chunk_size, n)
chunk = text[start:end]
chunks.append(chunk)
if end == n:
break
start = max(end - overlap, start + 1)
return chunks
def make_chunks(pages: List[Tuple[int, str]], source: str, chunk_size: int, overlap: int) -> List[Chunk]:
all_text = []
for page_num, raw in pages:
norm = normalize_text(raw)
if norm:
all_text.append((page_num, norm))
merged = "\n".join(t for _, t in all_text)
pieces = chunk_text(merged, chunk_size=chunk_size, overlap=overlap)
chunks = [Chunk(content=p, source=source) for p in pieces if p.strip()]
return chunks
def persist_chunks(session_id: str, chunks: List[Chunk]) -> str:
base = os.path.join(".tmp", "ingestion")
os.makedirs(base, exist_ok=True)
path = os.path.join(base, f"{session_id}.json")
with open(path, "w", encoding="utf-8") as f:
json.dump([asdict(c) for c in chunks], f, ensure_ascii=False, indent=2)
return path
def ingest_pdf(
file_path: str,
session_id: str,
parser: str = "pypdf",
chunk_size: int = 4000,
overlap: int = 400,
) -> Tuple[List[Chunk], str]:
if parser == "pypdf":
pages = extract_text_pypdf(file_path)
elif parser == "pdfplumber":
pages = extract_text_pdfplumber(file_path)
else:
raise ValueError("Unknown parser: " + parser)
chunks = make_chunks(pages, source=os.path.basename(file_path), chunk_size=chunk_size, overlap=overlap)
out = persist_chunks(session_id, chunks)
return chunks, out
|