| """Ingestion pipeline: parse PDFs -> chunk -> embed -> build & persist FAISS index. |
| |
| Run with: python -m src.ingest |
| |
| Chunks are built within page boundaries so every chunk maps to exactly one |
| page. This keeps citations exact and makes the eval's citation-validity |
| check programmatic. Pages longer than CHUNK_SIZE_TOKENS are split with a |
| sliding window of CHUNK_OVERLAP_TOKENS overlap. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
| import faiss |
| import numpy as np |
| from pypdf import PdfReader |
|
|
| import config |
|
|
| |
| |
| MIN_PAGE_CHARS = 25 |
|
|
|
|
| @dataclass |
| class Chunk: |
| """One retrievable unit of text, tied to a single page of one report.""" |
|
|
| id: int |
| company: str |
| doc: str |
| page: int |
| chunk_idx: int |
| text: str |
|
|
|
|
| def _extract_pages(pdf_path: Path) -> list[tuple[int, str]]: |
| """Extract (page_number, text) for every page; pdfplumber as fallback.""" |
| reader = PdfReader(pdf_path) |
| pages: list[tuple[int, str]] = [] |
| fallback_pages: list[int] = [] |
|
|
| for i, page in enumerate(reader.pages): |
| text = (page.extract_text() or "").strip() |
| if len(text) < MIN_PAGE_CHARS: |
| fallback_pages.append(i) |
| pages.append((i + 1, text)) |
|
|
| if fallback_pages: |
| import pdfplumber |
|
|
| with pdfplumber.open(pdf_path) as pdf: |
| for i in fallback_pages: |
| text = (pdf.pages[i].extract_text() or "").strip() |
| if len(text) >= MIN_PAGE_CHARS: |
| pages[i] = (i + 1, text) |
|
|
| return [(n, t) for n, t in pages if len(t) >= MIN_PAGE_CHARS] |
|
|
|
|
| def chunk_text(text: str, tokenizer, chunk_size: int, overlap: int) -> list[str]: |
| """Split text into chunks of ~chunk_size tokens with ~overlap token overlap. |
| |
| Token boundaries come from the embedding model's tokenizer so chunk sizes |
| match what the embedder actually sees. Chunks are sliced from the original |
| string via offset mappings (never decoded from token ids), preserving the |
| source casing and number formatting exactly. |
| """ |
| encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| offsets: list[tuple[int, int]] = encoding["offset_mapping"] |
| if len(offsets) <= chunk_size: |
| return [text] |
|
|
| chunks: list[str] = [] |
| step = chunk_size - overlap |
| start = 0 |
| while start < len(offsets): |
| window = offsets[start : start + chunk_size] |
| chunks.append(text[window[0][0] : window[-1][1]]) |
| if start + chunk_size >= len(offsets): |
| break |
| start += step |
| return chunks |
|
|
|
|
| def build_chunks(tokenizer) -> list[Chunk]: |
| """Parse every configured report into page-bound chunks.""" |
| chunks: list[Chunk] = [] |
| for stem, company in config.COMPANIES.items(): |
| pdf_path = config.REPORTS_DIR / f"{stem}.pdf" |
| if not pdf_path.exists(): |
| raise FileNotFoundError(f"Missing report: {pdf_path}") |
| pages = _extract_pages(pdf_path) |
| n_before = len(chunks) |
| for page_no, text in pages: |
| for idx, piece in enumerate( |
| chunk_text(text, tokenizer, config.CHUNK_SIZE_TOKENS, config.CHUNK_OVERLAP_TOKENS) |
| ): |
| chunks.append( |
| Chunk( |
| id=len(chunks), |
| company=company, |
| doc=stem, |
| page=page_no, |
| chunk_idx=idx, |
| text=piece, |
| ) |
| ) |
| print(f"{company}: {len(pages)} pages with text -> {len(chunks) - n_before} chunks") |
| return chunks |
|
|
|
|
| def main() -> None: |
| from sentence_transformers import SentenceTransformer |
|
|
| print(f"Loading embedding model {config.EMBEDDING_MODEL} ...") |
| model = SentenceTransformer(config.EMBEDDING_MODEL) |
| tokenizer = model.tokenizer |
|
|
| print("Parsing and chunking reports ...") |
| chunks = build_chunks(tokenizer) |
| print(f"Total: {len(chunks)} chunks") |
|
|
| print("Embedding chunks (normalized, cosine via inner product) ...") |
| embeddings = model.encode( |
| [c.text for c in chunks], |
| batch_size=64, |
| show_progress_bar=True, |
| normalize_embeddings=True, |
| ).astype(np.float32) |
|
|
| index = faiss.IndexFlatIP(embeddings.shape[1]) |
| index.add(embeddings) |
|
|
| config.INDEX_DIR.mkdir(parents=True, exist_ok=True) |
| faiss.write_index(index, str(config.INDEX_DIR / "index.faiss")) |
| with open(config.INDEX_DIR / "chunks.jsonl", "w", encoding="utf-8") as f: |
| for c in chunks: |
| f.write(json.dumps(asdict(c), ensure_ascii=False) + "\n") |
|
|
| print(f"Wrote {config.INDEX_DIR / 'index.faiss'} ({index.ntotal} vectors) " |
| f"and chunks.jsonl") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
