Update README.md

README.md

@@ -17,7 +17,7 @@ license: apache-2.0
 <h1>DocExplainerV0: Visual Document QA with Bounding Box Localization</h1>
 
 [![License: CC BY 4.0](https://img.shields.io/badge/License-CC%20BY%204.0-lightgrey.svg)](https://creativecommons.org/licenses/by/4.0/)
-[![arXiv](https://img.shields.io/badge/arXiv-2501.03403-b31b1b.svg)]()
+<!-- [![arXiv](https://img.shields.io/badge/arXiv-2501.03403-b31b1b.svg)]() -->
 [![HuggingFace](https://img.shields.io/badge/🤗%20Hugging%20Face-Datasets-yellow)](https://huggingface.co/letxbe/DocExplainerV0)
 
 </div>
@@ -78,17 +78,73 @@ Here is a simple example of how to use `DocExplainer` to get an answer and its c
 ```python
 from PIL import Image
 import requests
-from transformers import AutoModel
+import torch
+from transformers import AutoModel, AutoModelForImageTextToText, AutoProcessor
+import json 
 
-# Load example document image
-url = "https://datasets-server.huggingface.co/cached-assets/letxbe/BoundingDocs/--/47db6d2b6af0aadfd082591a8445d0f47c3b8d61/--/default/test/7/doc_images/image-1d100e9.jpg"
+url = "https://i.postimg.cc/BvftyvS3/image-1d100e9.jpg"
 image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
 question = "What is the invoice number?"
-answer = "3Y8M2d-846" # generate it with any VLM
 
-explainer = AutoModel.from_pretrained("letxbe/DocExplainerv0", trust_remote_code=True)
+# -----------------------
+# 1. Load SmolVLM2-2.2B for answer generation
+# -----------------------
+vlm_model = AutoModelForImageTextToText.from_pretrained(
+    "HuggingFaceTB/SmolVLM2-2.2B-Instruct",
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    attn_implementation="flash_attention_2"
+)
+processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct")
+
+PROMPT = """Based only on the document image, answer the following question:
+Question: {QUESTION}
+Provide ONLY a JSON response in the following format (no trailing commas!):
+{{
+  "content": "answer"
+}}
+"""
+
+prompt_text = PROMPT.format(QUESTION=question)
+
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": prompt_text},
+        ]
+    },
+]
+
+inputs = processor.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    tokenize=True,
+    return_dict=True,
+    return_tensors="pt",
+).to(vlm_model.device, dtype=torch.bfloat16)
+
+input_length = inputs['input_ids'].shape[1]
+generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=2056)
+
+output_ids = generated_ids[:, input_length:]
+generated_texts = processor.batch_decode(
+    output_ids,
+    skip_special_tokens=True,
+)
+
+decoded_output = generated_texts[0].replace("Assistant:", "", 1).strip()
+answer = json.loads(decoded_output)['content']
+
+print(f"Answer: {answer}")
+
+# -----------------------
+# 2. Load DocExplainerV0 for bounding box prediction
+# -----------------------
+explainer = AutoModel.from_pretrained("letxbe/DocExplainerV0", trust_remote_code=True)
 bbox = explainer.predict(image, answer)
-print(f"Bounding box: {bbox}")  # [x1, y1, x2, y2]
+print(f"Predicted bounding box (normalized): {bbox}")
 ```
 
 
@@ -109,30 +165,6 @@ Example Output:
 </table> 
 
 
-
-
-## Performance 
-Evaluated on [BoundingDocs v2.0](https://huggingface.co/datasets/letxbe/BoundingDocs) dataset: 
-### Full DocExplainer Pipeline 
-
-| VLM Model       | ANLS ↑| IoU ↑ |
-| --------------- | ----- | ----- | 
-| SmolVLM2-2.2b   | 0.572 | 0.175 |
-| qwen2.5-vl-7b   | 0.689 | 0.188 |
-
-     
-### VLM-only Baseline (for comparison)
-| VLM Model       | ANLS ↑| IoU ↑ |
-| --------------- | ----- | ----- | 
-| SmolVLM2-2.2b   | 0.561 | 0.011 |
-| qwen2.5-vl-7b   | 0.720 | 0.038 |
-| Claude Sonnet 4 | 0.737 | 0.031 |
-
-
-
-
-
-
 ## Limitations
 - **Prototype only**: Intended as a first approach, not a production-ready solution.
 - **Dataset constraints**: Current evaluation is limited to cases where an answer fits in a single bounding box. Answers requiring reasoning over multiple regions or not fully captured by OCR cannot be properly evaluated.