kostissz commited on
Commit
37e70be
·
verified ·
1 Parent(s): 2e8fc61

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -2
app.py CHANGED
@@ -1,15 +1,22 @@
1
  from typing import Dict, Tuple
2
  import os
3
  import gradio as gr
 
4
  from docling.datamodel.base_models import InputFormat
5
- from docling.datamodel.pipeline_options import PdfPipelineOptions
6
  from docling.document_converter import DocumentConverter, PdfFormatOption
7
  from docling_core.types import DoclingDocument
8
  from docling.utils import model_downloader
9
  from docling.datamodel.pipeline_options import smolvlm_picture_description
10
 
11
  # Download models upon HF space initialization
 
 
 
 
 
12
  if os.getenv("IS_HF_SPACE"):
 
13
  model_downloader.download_models()
14
 
15
 
@@ -22,14 +29,17 @@ def parse_document(
22
  ) -> Tuple[DoclingDocument, str]:
23
  yield None, f"Parsing document... ⏳"
24
 
25
- pipeline_options = PdfPipelineOptions()
26
  pipeline_options.do_code_enrichment = do_code_enrichment
27
  pipeline_options.do_formula_enrichment = do_formula_enrichment
28
  pipeline_options.generate_picture_images = do_picture_classification
29
  pipeline_options.images_scale = 2
30
  pipeline_options.do_picture_classification = do_picture_classification
 
31
  pipeline_options.do_picture_description = do_picture_description
32
  pipeline_options.picture_description_options = smolvlm_picture_description
 
 
 
33
 
34
  print(f"Pipeline options defined: \n\t{pipeline_options}")
35
  converter = DocumentConverter(
 
1
  from typing import Dict, Tuple
2
  import os
3
  import gradio as gr
4
+ import torch.cuda
5
  from docling.datamodel.base_models import InputFormat
6
+ from docling.datamodel.pipeline_options import PdfPipelineOptions, AcceleratorDevice
7
  from docling.document_converter import DocumentConverter, PdfFormatOption
8
  from docling_core.types import DoclingDocument
9
  from docling.utils import model_downloader
10
  from docling.datamodel.pipeline_options import smolvlm_picture_description
11
 
12
  # Download models upon HF space initialization
13
+ pipeline_options = PdfPipelineOptions()
14
+ if torch.cuda.is_available():
15
+ print("Enabling CUDA Accelerator")
16
+ pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA
17
+ pipeline_options.accelerator_options.cuda_use_flash_attention2 = True
18
  if os.getenv("IS_HF_SPACE"):
19
+ print("Downloading models...")
20
  model_downloader.download_models()
21
 
22
 
 
29
  ) -> Tuple[DoclingDocument, str]:
30
  yield None, f"Parsing document... ⏳"
31
 
 
32
  pipeline_options.do_code_enrichment = do_code_enrichment
33
  pipeline_options.do_formula_enrichment = do_formula_enrichment
34
  pipeline_options.generate_picture_images = do_picture_classification
35
  pipeline_options.images_scale = 2
36
  pipeline_options.do_picture_classification = do_picture_classification
37
+
38
  pipeline_options.do_picture_description = do_picture_description
39
  pipeline_options.picture_description_options = smolvlm_picture_description
40
+ pipeline_options.picture_description_options.prompt = "Describe the image in three sentences. Be concise and accurate."
41
+ pipeline_options.images_scale = 2.0
42
+ pipeline_options.generate_picture_images = True
43
 
44
  print(f"Pipeline options defined: \n\t{pipeline_options}")
45
  converter = DocumentConverter(