langchain/langchain/document_loaders/paged_pdf.py
Harrison Chase c64f98e2bb
Harrison/format agent instructions (#973)
Co-authored-by: Andrew White <white.d.andrew@gmail.com>
Co-authored-by: Harrison Chase <harrisonchase@Harrisons-MBP.attlocal.net>
Co-authored-by: Peng Qu <82029664+pengqu123@users.noreply.github.com>
2023-02-10 10:07:26 -08:00

37 lines
1.1 KiB
Python

"""Loads a PDF with pypdf and chunks at character level."""
from typing import List
from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader
class PagedPDFSplitter(BaseLoader):
"""Loads a PDF with pypdf and chunks at character level.
Loader also stores page numbers in metadatas.
"""
def __init__(self, file_path: str):
"""Initialize with file path."""
try:
import pypdf # noqa:F401
except ImportError:
raise ValueError(
"pypdf package not found, please install it with " "`pip install pypdf`"
)
self._file_path = file_path
def load(self) -> List[Document]:
"""Load given path as pages."""
import pypdf
pdf_file_obj = open(self._file_path, "rb")
pdf_reader = pypdf.PdfReader(pdf_file_obj)
docs = []
for i, page in enumerate(pdf_reader.pages):
text = page.extract_text()
metadata = {"source": self._file_path, "page": i}
docs.append(Document(page_content=text, metadata=metadata))
pdf_file_obj.close()
return docs