Spaces:
Running
Running
File size: 1,373 Bytes
170d9a9 |
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 |
import asyncio
from medrag_multi_modal.document_loader import (
PDFPlumberTextLoader,
PyMuPDF4LLMTextLoader,
PyPDF2TextLoader,
)
URL = "https://archive.org/download/GraysAnatomy41E2015PDF/Grays%20Anatomy-41%20E%20%282015%29%20%5BPDF%5D.pdf"
COLUMN_NAMES = [
"text",
"page_idx",
"document_name",
"file_path",
"file_url",
"loader_name",
]
def test_pdfplumber_text_loader():
loader = PDFPlumberTextLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=31, end_page=36))
assert dataset.num_rows == 6
assert dataset.column_names == COLUMN_NAMES
def test_pymupdf_text_loader():
loader = PyMuPDF4LLMTextLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=31, end_page=36))
assert dataset.num_rows == 6
assert dataset.column_names == COLUMN_NAMES
def test_pypdf2_text_loader():
loader = PyPDF2TextLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=31, end_page=36))
assert dataset.num_rows == 6
assert dataset.column_names == COLUMN_NAMES
|