taprosoft
commited on
Commit
·
3b036fa
1
Parent(s):
9f1821c
fix: update flash_attn
Browse files- app.py +6 -0
- backends/smoldocling.py +1 -1
- 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="
|
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
|
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
|