taprosoft commited on
Commit
3b036fa
·
1 Parent(s): 9f1821c

fix: update flash_attn

Browse files
Files changed (3) hide show
  1. app.py +6 -0
  2. backends/smoldocling.py +1 -1
  3. requirements.txt +1 -1
app.py CHANGED
@@ -28,9 +28,15 @@ MAX_PAGES = int(os.getenv("MAX_PAGES", "2"))
28
 
29
  # Install poppler-utils
30
  import os
 
31
 
32
  os.system("apt-get update")
33
  os.system("apt-get install --no-install-recommends -y poppler-utils tesseract-ocr")
 
 
 
 
 
34
 
35
 
36
  def convert_document(path, method, start_page=0, enabled=True):
 
28
 
29
  # Install poppler-utils
30
  import os
31
+ import subprocess
32
 
33
  os.system("apt-get update")
34
  os.system("apt-get install --no-install-recommends -y poppler-utils tesseract-ocr")
35
+ subprocess.run(
36
+ "pip3 install flash-attn --no-build-isolation",
37
+ env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
38
+ shell=True,
39
+ )
40
 
41
 
42
  def convert_document(path, method, start_page=0, enabled=True):
backends/smoldocling.py CHANGED
@@ -19,7 +19,7 @@ processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
19
  model = AutoModelForVision2Seq.from_pretrained(
20
  "ds4sd/SmolDocling-256M-preview",
21
  torch_dtype=torch.bfloat16,
22
- _attn_implementation="sdpa",
23
  ).to(DEVICE)
24
 
25
  # Create input messages
 
19
  model = AutoModelForVision2Seq.from_pretrained(
20
  "ds4sd/SmolDocling-256M-preview",
21
  torch_dtype=torch.bfloat16,
22
+ _attn_implementation="flash_attention_2",
23
  ).to(DEVICE)
24
 
25
  # Create input messages
requirements.txt CHANGED
@@ -5,7 +5,7 @@ unstructured[pdf]
5
  openai
6
  img2table
7
  gmft
8
- transformers<5.0.0,>=4.45.2
9
  pypdf
10
  docling_core
11
  opencv-contrib-python
 
5
  openai
6
  img2table
7
  gmft
8
+ transformers
9
  pypdf
10
  docling_core
11
  opencv-contrib-python