Update README.md

README.md

@@ -95,7 +95,8 @@ doc.load_from_document_tokens([doctags], [image])
 # export as any format
 # HTML
 # print(doc.export_to_html())
-
+# with open(output_file, "w", encoding="utf-8") as f:
+#     f.write(doc.export_to_html())
 # MD
 # print(doc.export_to_markdown())
 
@@ -115,7 +116,8 @@ from transformers.image_utils import load_image
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load images
-image = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")
+page_1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")
+page_2 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")
 
 # Initialize processor and model
 processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
@@ -131,24 +133,39 @@ messages = [
         "role": "user",
         "content": [
             {"type": "image"},
-            {"type": "text", "text": "Convert this page to docling."}
+            {"type": "image"},
+            {"type": "text", "text": "Convert this document to docling."}
         ]
     },
 ]
 
 # Prepare inputs
 prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-inputs = processor(text=prompt, images=[image], return_tensors="pt")
+inputs = processor(text=prompt, images=[page_1, page_2], return_tensors="pt")
 inputs = inputs.to(DEVICE)
 
 # Generate outputs
-generated_ids = model.generate(**inputs, max_new_tokens=500)
-generated_texts = processor.batch_decode(
-    generated_ids,
-    skip_special_tokens=True,
-)
+generated_ids = model.generate(**inputs, max_new_tokens=8192)
+prompt_length = inputs.input_ids.shape[1]
+trimmed_generated_ids = generated_ids[:, prompt_length:]
+doctags = processor.batch_decode(
+    trimmed_generated_ids,
+    skip_special_tokens=False,
+)[0].lstrip()
+
+# create a docling document
+doc = DoclingDocument(name="Document")
+
+# populate it
+doc.load_from_document_tokens([doctags], [page_1, page_2])
 
-print(generated_texts[0])
+# export as any format
+# HTML
+# print(doc.export_to_html())
+# with open(output_file, "w", encoding="utf-8") as f:
+#     f.write(doc.export_to_html())
+# MD
+# print(doc.export_to_markdown())
 ``````
 </details>