Update README.md

README.md

@@ -101,24 +101,25 @@ You can use **transformers**, **vllm**, or **onnx** to perform inference, and [D
 # pip install torch
 # pip install docling_core
 # pip install transformers
+# pip install flash_attn --no-build-isolation # on supported Nvidia GPUs
 
 import torch
 from docling_core.types.doc import DoclingDocument
 from docling_core.types.doc.document import DocTagsDocument
-from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers import AutoProcessor, AutoModelForImageTextToText
 from transformers.image_utils import load_image
 from pathlib import Path
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load images
-image = load_image("https://upload.wikimedia.org/wikipedia/commons/7/76/GazettedeFrance.jpg")
+image = load_image("https://huggingface.co/ibm-granite/granite-docling-258M/resolve/main/assets/new_arxiv.png")
 
 # Initialize processor and model
 processor = AutoProcessor.from_pretrained("ibm-granite/granite-docling-258M")
-model = AutoModelForVision2Seq.from_pretrained(
+model = AutoModelForImageTextToText.from_pretrained(
     "ibm-granite/granite-docling-258M",
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "sdpa",
 ).to(DEVICE)
 
@@ -128,7 +129,7 @@ messages = [
         "role": "user",
         "content": [
             {"type": "image"},
-            {"type": "text", "text": "Convert this page to docling."}
+            {"type": "text", "text": "Convert page to docling."}
         ]
     },
 ]