| 33 | |
| 34 | |
| 35 | class DocLoader: |
| 36 | |
| 37 | MAX_NUM_PAGES = 1000 |
| 38 | MAX_FILE_SIZE = 100 * 1024 * 1024 # 100MB |
| 39 | |
| 40 | # Number of threads for document conversion |
| 41 | DOC_CONVERT_NUM_THREADS = os.environ.get('DOC_CONVERT_NUM_THREADS', 16) |
| 42 | |
| 43 | def __init__(self, verbose: bool = False): |
| 44 | |
| 45 | self._verbose = verbose |
| 46 | |
| 47 | accelerator_options = AcceleratorOptions() |
| 48 | accelerator_options.num_threads = self.DOC_CONVERT_NUM_THREADS |
| 49 | accelerator_options.device = 'auto' |
| 50 | accelerator_options.cuda_use_flash_attention2 = False # TODO |
| 51 | |
| 52 | pdf_pipeline_options = PdfPipelineOptions() |
| 53 | pdf_pipeline_options.generate_page_images = True |
| 54 | pdf_pipeline_options.generate_picture_images = True |
| 55 | pdf_pipeline_options.generate_table_images = True |
| 56 | pdf_pipeline_options.do_code_enrichment = False |
| 57 | pdf_pipeline_options.do_formula_enrichment = False |
| 58 | pdf_pipeline_options.do_picture_classification = True |
| 59 | pdf_pipeline_options.do_picture_description = False |
| 60 | pdf_pipeline_options.images_scale = 2.0 |
| 61 | pdf_pipeline_options.accelerator_options = accelerator_options # type: ignore |
| 62 | |
| 63 | self._converter = DocumentConverter( |
| 64 | format_options={ |
| 65 | InputFormat.PDF: |
| 66 | PdfFormatOption(pipeline_options=pdf_pipeline_options) |
| 67 | }) |
| 68 | |
| 69 | @staticmethod |
| 70 | def _group_by_input_format( |
| 71 | urls_or_files: list[str]) -> dict[InputFormat, list[str]]: |
| 72 | """ |
| 73 | Group the provided URLs or files by their input format. |
| 74 | This is a placeholder implementation and should be replaced with actual logic. |
| 75 | |
| 76 | TODO: to be implemented in the future, currently only supports PDF and HTML formats. |
| 77 | """ |
| 78 | grouped = {InputFormat.PDF: [], InputFormat.HTML: []} |
| 79 | for url_or_file in urls_or_files: |
| 80 | if url_or_file.endswith('.pdf') or url_or_file.startswith( |
| 81 | 'https://arxiv.org/pdf/'): |
| 82 | grouped[InputFormat.PDF].append(url_or_file) |
| 83 | elif url_or_file.startswith('http') or url_or_file.endswith( |
| 84 | '.html'): |
| 85 | grouped[InputFormat.HTML].append(url_or_file) |
| 86 | else: |
| 87 | logger.error( |
| 88 | f'**Error: Unsupported file type for {url_or_file}') |
| 89 | return grouped |
| 90 | |
| 91 | @staticmethod |
| 92 | def _transform_dict(original_dict): |
no test coverage detected