# pyright: ignore
from docling.datamodel.accelerator_options import AcceleratorDevice
from docling.datamodel.base_models import InputFormat

# Create converter with default VLM (Granite-Docling)
from docling.datamodel.pipeline_options import VlmPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.vlm_pipeline import VlmPipeline

pipeline_options = VlmPipelineOptions()
pipeline_options.generate_page_images = True  # Enable efficient page batching
pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA
# pipeline_options.accelerator_options.cuda_use_flash_attention2 = True  # Comment this out

converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(
            pipeline_cls=VlmPipeline,
            pipeline_options=pipeline_options,
        ),
    }
)

# Convert your PDF
pdf_path = "documents/TMUS_992_Q423.pdf"  # Change this to your PDF path
result = converter.convert(pdf_path)
document = result.document

# Save as markdown
output_path = "converted_document.md"
with open(output_path, "w") as f:
    f.write(document.export_to_markdown())

print(f"Converted {pdf_path} to {output_path}")
print(f"Found {len(document.tables)} tables")