Spaces:
Running
Running
| import os | |
| import tempfile | |
| from typing import List, Dict, Any, BinaryIO | |
| import pdfplumber | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from config import CHUNK_SIZE, CHUNK_OVERLAP | |
| """ | |
| Methods: | |
| extract_uploaded_pdf_pages -> Takes in a file_obj in the form of BinaryIO and a filename as string. Extracts text from PDF file, page by page. | |
| extract_pdf_from_path -> Extract text from PDF file path directly for | |
| chunk_text -> Split text into chunks with metadata. | |
| process_documents -> Process extracted pages into chunks. | |
| chunks_to_store_format -> Convert chunks to vector store add_documents format. | |
| """ | |
| def extract_pdf_pages(file_obj: BinaryIO, filename: str = "document.pdf") -> List[Dict[str, Any]]: | |
| pages = [] | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: | |
| tmp.write(file_obj.read()) | |
| tmp_path = tmp.name | |
| try: | |
| with pdfplumber.open(tmp_path) as pdf: | |
| for i, page in enumerate(pdf.pages, 1): | |
| text = page.extract_text() or "" | |
| if text.strip(): | |
| pages.append({ | |
| "text": text, | |
| "page_number": i, | |
| "source": filename | |
| }) | |
| finally: | |
| os.unlink(tmp_path) | |
| return pages | |
| def extract_pdf_from_path(pdf_path: str, filename: str = None) -> List[Dict[str, Any]]: | |
| pages = [] | |
| if filename is None: | |
| filename = os.path.basename(pdf_path) | |
| try: | |
| with pdfplumber.open(pdf_path) as pdf: | |
| for i, page in enumerate(pdf.pages, 1): | |
| text = page.extract_text() or "" | |
| if text.strip(): | |
| pages.append({ | |
| "text": text, | |
| "page_number": i, | |
| "source": filename | |
| }) | |
| except Exception: | |
| pass | |
| return pages | |
| def chunk_text(text: str, source: str, page_number: int, chunk_size: int = CHUNK_SIZE, chunk_overlap: int = CHUNK_OVERLAP) -> List[Dict[str, Any]]: | |
| splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=chunk_size, | |
| chunk_overlap=chunk_overlap, | |
| separators=["\n\n", "\n", ".", " ", ""] | |
| ) | |
| chunks = splitter.split_text(text) | |
| result = [] | |
| for i, chunk in enumerate(chunks): | |
| result.append({ | |
| "text": chunk, | |
| "metadata": { | |
| "source": source, | |
| "page_number": page_number, | |
| "chunk_index": i | |
| } | |
| }) | |
| return result | |
| def process_documents(pages: List[Dict[str, Any]], chunk_size: int = CHUNK_SIZE, chunk_overlap: int = CHUNK_OVERLAP) -> List[Dict[str, Any]]: | |
| all_chunks = [] | |
| for page in pages: | |
| chunks = chunk_text( | |
| text=page["text"], | |
| source=page["source"], | |
| page_number=page["page_number"], | |
| chunk_size=chunk_size, | |
| chunk_overlap=chunk_overlap | |
| ) | |
| all_chunks.extend(chunks) | |
| return all_chunks | |
| def chunks_to_store_format(chunks: List[Dict[str, Any]]) -> tuple: | |
| texts = [c["text"] for c in chunks] | |
| metadatas = [c["metadata"] for c in chunks] | |
| return texts, metadatas | |
| if __name__ == "__main__": | |
| test_text = """This is a test document. | |
| It has multiple paragraphs. Each paragraph contains some text. | |
| This is the third paragraph with more content.""" | |
| chunks = chunk_text(test_text, "test.pdf", 1) | |
| print(f"Created {len(chunks)} chunks") | |
| for c in chunks: | |
| print(f" - {c['metadata']}: {c['text'][:50]}...") | |