{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/docling-project/docling/blob/main/docs/examples/visual_grounding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visual grounding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Step | Tech | Execution | \n",
    "| --- | --- | --- |\n",
    "| Embedding | Hugging Face / Sentence Transformers | 💻 Local |\n",
    "| Vector store | Milvus | 💻 Local |\n",
    "| Gen AI | Hugging Face Inference API | 🌐 Remote | "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example showcases Docling's **visual grounding** capabilities, which can be combined\n",
    "with any agentic AI / RAG framework.\n",
    "\n",
    "In this instance, we illustrate these capabilities leveraging the\n",
    "[LangChain Docling integration](../../integrations/langchain/), along with a Milvus\n",
    "vector store, as well as sentence-transformers embeddings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 👉 For best conversion speed, use GPU acceleration whenever available; e.g. if running on Colab, use GPU-enabled runtime.\n",
    "- Notebook uses HuggingFace's Inference API; for increased LLM quota, token can be provided via env var `HF_TOKEN`.\n",
    "- Requirements can be installed as shown below (`--no-warn-conflicts` meant for Colab's pre-populated Python env; feel free to remove for stricter usage):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q --progress-bar off --no-warn-conflicts langchain-docling langchain-core langchain-huggingface langchain_milvus langchain matplotlib python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "from tempfile import mkdtemp\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_docling.loader import ExportType\n",
    "\n",
    "\n",
    "def _get_env_from_colab_or_os(key):\n",
    "    try:\n",
    "        from google.colab import userdata\n",
    "\n",
    "        try:\n",
    "            return userdata.get(key)\n",
    "        except userdata.SecretNotFoundError:\n",
    "            pass\n",
    "    except ImportError:\n",
    "        pass\n",
    "    return os.getenv(key)\n",
    "\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "# https://github.com/huggingface/transformers/issues/5486:\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "HF_TOKEN = _get_env_from_colab_or_os(\"HF_TOKEN\")\n",
    "SOURCES = [\"https://arxiv.org/pdf/2408.09869\"]  # Docling Technical Report\n",
    "EMBED_MODEL_ID = \"sentence-transformers/all-MiniLM-L6-v2\"\n",
    "GEN_MODEL_ID = \"mistralai/Mixtral-8x7B-Instruct-v0.1\"\n",
    "QUESTION = \"Which are the main AI models in Docling?\"\n",
    "PROMPT = PromptTemplate.from_template(\n",
    "    \"Context information is below.\\n---------------------\\n{context}\\n---------------------\\nGiven the context information and not prior knowledge, answer the query.\\nQuery: {input}\\nAnswer:\\n\",\n",
    ")\n",
    "TOP_K = 3\n",
    "MILVUS_URI = str(Path(mkdtemp()) / \"docling.db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Document store setup\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Document loading\n",
    "\n",
    "We first define our converter, in this case including options for keeping page images (for visual grounding)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from docling.datamodel.base_models import InputFormat\n",
    "from docling.datamodel.pipeline_options import PdfPipelineOptions\n",
    "from docling.document_converter import DocumentConverter, PdfFormatOption\n",
    "\n",
    "converter = DocumentConverter(\n",
    "    format_options={\n",
    "        InputFormat.PDF: PdfFormatOption(\n",
    "            pipeline_options=PdfPipelineOptions(\n",
    "                generate_page_images=True,\n",
    "                images_scale=2.0,\n",
    "            ),\n",
    "        )\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set up a simple doc store for keeping converted documents, as that is needed for visual grounding further below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_store = {}\n",
    "doc_store_root = Path(mkdtemp())\n",
    "for source in SOURCES:\n",
    "    dl_doc = converter.convert(source=source).document\n",
    "    file_path = Path(doc_store_root / f\"{dl_doc.origin.binary_hash}.json\")\n",
    "    dl_doc.save_as_json(file_path)\n",
    "    doc_store[dl_doc.origin.binary_hash] = file_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can instantiate our loader and load documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (648 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    }
   ],
   "source": [
    "from langchain_docling import DoclingLoader\n",
    "\n",
    "from docling.chunking import HybridChunker\n",
    "\n",
    "loader = DoclingLoader(\n",
    "    file_path=SOURCES,\n",
    "    converter=converter,\n",
    "    export_type=ExportType.DOC_CHUNKS,\n",
    "    chunker=HybridChunker(tokenizer=EMBED_MODEL_ID),\n",
    ")\n",
    "\n",
    "docs = loader.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 👉 **NOTE**: As you see above, using the `HybridChunker` can sometimes lead to a warning from the transformers library, however this is a \"false alarm\" — for details check [here](https://docling-project.github.io/docling/faq/#hybridchunker-triggers-warning-token-indices-sequence-length-is-longer-than-the-specified-maximum-sequence-length-for-this-model)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspecting some sample splits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- d.page_content='Docling Technical Report\\nVersion 1.0\\nChristoph Auer Maksym Lysak Ahmed Nassar Michele Dolfi Nikolaos Livathinos Panos Vagenas Cesar Berrospi Ramis Matteo Omenetti Fabian Lindlbauer Kasper Dinkla Lokesh Mishra Yusik Kim Shubham Gupta Rafael Teixeira de Lima Valery Weber Lucas Morin Ingmar Meijer Viktor Kuropiatnyk Peter W. J. Staar\\nAI4K Group, IBM Research R¨ uschlikon, Switzerland'\n",
      "- d.page_content='Abstract\\nThis technical report introduces Docling , an easy to use, self-contained, MITlicensed open-source package for PDF document conversion. It is powered by state-of-the-art specialized AI models for layout analysis (DocLayNet) and table structure recognition (TableFormer), and runs efficiently on commodity hardware in a small resource budget. The code interface allows for easy extensibility and addition of new features and models.'\n",
      "- d.page_content='1 Introduction\\nConverting PDF documents back into a machine-processable format has been a major challenge for decades due to their huge variability in formats, weak standardization and printing-optimized characteristic, which discards most structural features and metadata. With the advent of LLMs and popular application patterns such as retrieval-augmented generation (RAG), leveraging the rich content embedded in PDFs has become ever more relevant. In the past decade, several powerful document understanding solutions have emerged on the market, most of which are commercial software, cloud offerings [3] and most recently, multi-modal vision-language models. As of today, only a handful of open-source tools cover PDF conversion, leaving a significant feature and quality gap to proprietary solutions.\\nWith Docling , we open-source a very capable and efficient document conversion tool which builds on the powerful, specialized AI models and datasets for layout analysis and table structure recognition we developed and presented in the recent past [12, 13, 9]. Docling is designed as a simple, self-contained python library with permissive license, running entirely locally on commodity hardware. Its code architecture allows for easy extensibility and addition of new features and models.\\nHere is what Docling delivers today:\\n· Converts PDF documents to JSON or Markdown format, stable and lightning fast\\n· Understands detailed page layout, reading order, locates figures and recovers table structures\\n· Extracts metadata from the document, such as title, authors, references and language\\n· Optionally applies OCR, e.g. for scanned PDFs\\n· Can be configured to be optimal for batch-mode (i.e high throughput, low time-to-solution) or interactive mode (compromise on efficiency, low time-to-solution)\\n· Can leverage different accelerators (GPU, MPS, etc).'\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "for d in docs[:3]:\n",
    "    print(f\"- {d.page_content=}\")\n",
    "print(\"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ingestion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "from tempfile import mkdtemp\n",
    "\n",
    "from langchain_huggingface.embeddings import HuggingFaceEmbeddings\n",
    "from langchain_milvus import Milvus\n",
    "\n",
    "embedding = HuggingFaceEmbeddings(model_name=EMBED_MODEL_ID)\n",
    "\n",
    "\n",
    "milvus_uri = str(Path(mkdtemp()) / \"docling.db\")  # or set as needed\n",
    "vectorstore = Milvus.from_documents(\n",
    "    documents=docs,\n",
    "    embedding=embedding,\n",
    "    collection_name=\"docling_demo\",\n",
    "    connection_args={\"uri\": milvus_uri},\n",
    "    index_params={\"index_type\": \"FLAT\"},\n",
    "    drop_old=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured.\n"
     ]
    }
   ],
   "source": [
    "from langchain.chains import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "from langchain_huggingface import HuggingFaceEndpoint\n",
    "\n",
    "retriever = vectorstore.as_retriever(search_kwargs={\"k\": TOP_K})\n",
    "llm = HuggingFaceEndpoint(\n",
    "    repo_id=GEN_MODEL_ID,\n",
    "    huggingfacehub_api_token=HF_TOKEN,\n",
    ")\n",
    "\n",
    "\n",
    "def clip_text(text, threshold=100):\n",
    "    return f\"{text[:threshold]}...\" if len(text) > threshold else text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/pva/work/github.com/DS4SD/docling/.venv/lib/python3.12/site-packages/huggingface_hub/utils/_deprecation.py:131: FutureWarning: 'post' (from 'huggingface_hub.inference._client') is deprecated and will be removed from version '0.31.0'. Making direct POST requests to the inference server is not supported anymore. Please use task methods instead (e.g. `InferenceClient.chat_completion`). If your use case is not supported, please open an issue in https://github.com/huggingface/huggingface_hub.\n",
      "  warnings.warn(warning_message, FutureWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "Which are the main AI models in Docling?\n",
      "\n",
      "Answer:\n",
      "The main AI models in Docling are:\n",
      "1. A layout analysis model, an accurate object-detector for page elements.\n",
      "2. TableFormer, a state-of-the-art table structure recognition model.\n"
     ]
    }
   ],
   "source": [
    "from docling.chunking import DocMeta\n",
    "from docling.datamodel.document import DoclingDocument\n",
    "\n",
    "question_answer_chain = create_stuff_documents_chain(llm, PROMPT)\n",
    "rag_chain = create_retrieval_chain(retriever, question_answer_chain)\n",
    "resp_dict = rag_chain.invoke({\"input\": QUESTION})\n",
    "\n",
    "clipped_answer = clip_text(resp_dict[\"answer\"], threshold=200)\n",
    "print(f\"Question:\\n{resp_dict['input']}\\n\\nAnswer:\\n{clipped_answer}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual grounding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source 1:\n",
      "  text: \"3.2 AI models\\nAs part of Docling, we initially release two highly capable AI models to the open-source community, which have been developed and published recently by our team. The first model is a layout analysis model, an accurate object-detector for page elements [13]. The second model is TableFormer [12, 9], a state-of-the-art table structure re...\"\n",
      "  page: 3\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 407980 characters]",
      "text/plain": [
       "<Figure size 1500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source 2:\n",
      "  text: \"3 Processing pipeline\\nDocling implements a linear pipeline of operations, which execute sequentially on each given document (see Fig. 1). Each document is first parsed by a PDF backend, which retrieves the programmatic text tokens, consisting of string content and its coordinates on the page, and also renders a bitmap image of each page to support ...\"\n",
      "  page: 2\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 423756 characters]",
      "text/plain": [
       "<Figure size 1500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source 3:\n",
      "  text: \"6 Future work and contributions\\nDocling is designed to allow easy extension of the model library and pipelines. In the future, we plan to extend Docling with several more models, such as a figure-classifier model, an equationrecognition model, a code-recognition model and more. This will help improve the quality of conversion for specific types of ...\"\n",
      "  page: 5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 444596 characters]",
      "text/plain": [
       "<Figure size 1500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from PIL import ImageDraw\n",
    "\n",
    "for i, doc in enumerate(resp_dict[\"context\"][:]):\n",
    "    image_by_page = {}\n",
    "    print(f\"Source {i + 1}:\")\n",
    "    print(f\"  text: {json.dumps(clip_text(doc.page_content, threshold=350))}\")\n",
    "    meta = DocMeta.model_validate(doc.metadata[\"dl_meta\"])\n",
    "\n",
    "    # loading the full DoclingDocument from the document store:\n",
    "    dl_doc = DoclingDocument.load_from_json(doc_store.get(meta.origin.binary_hash))\n",
    "\n",
    "    for doc_item in meta.doc_items:\n",
    "        if doc_item.prov:\n",
    "            prov = doc_item.prov[0]  # here we only consider the first provenence item\n",
    "            page_no = prov.page_no\n",
    "            if img := image_by_page.get(page_no):\n",
    "                pass\n",
    "            else:\n",
    "                page = dl_doc.pages[prov.page_no]\n",
    "                print(f\"  page: {prov.page_no}\")\n",
    "                img = page.image.pil_image\n",
    "                image_by_page[page_no] = img\n",
    "            bbox = prov.bbox.to_top_left_origin(page_height=page.size.height)\n",
    "            bbox = bbox.normalized(page.size)\n",
    "            thickness = 2\n",
    "            padding = thickness + 2\n",
    "            bbox.l = round(bbox.l * img.width - padding)\n",
    "            bbox.r = round(bbox.r * img.width + padding)\n",
    "            bbox.t = round(bbox.t * img.height - padding)\n",
    "            bbox.b = round(bbox.b * img.height + padding)\n",
    "            draw = ImageDraw.Draw(img)\n",
    "            draw.rectangle(\n",
    "                xy=bbox.as_tuple(),\n",
    "                outline=\"blue\",\n",
    "                width=thickness,\n",
    "            )\n",
    "    for p in image_by_page:\n",
    "        img = image_by_page[p]\n",
    "        plt.figure(figsize=[15, 15])\n",
    "        plt.imshow(img)\n",
    "        plt.axis(\"off\")\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
