Demo_ocrtest_qwen.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **FT; Key Information Extraction**\n",
        "Qwen2VLForConditionalGeneration"
      ],
      "metadata": {
        "id": "-b4-SW1aGOcF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install gradio spaces transformers accelerate numpy requests torch torchvision qwen-vl-utils av ipython reportlab fpdf python-docx pillow huggingface_hub"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oDmd1ZObGSel",
        "outputId": "5b01f267-d5af-4409-cf67-6c318388d584"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting gradio\n",
            "  Downloading gradio-5.9.1-py3-none-any.whl.metadata (16 kB)\n",
            "Collecting spaces\n",
            "  Downloading spaces-0.31.1-py3-none-any.whl.metadata (1.0 kB)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.47.1)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (1.2.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (1.26.4)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (2.32.3)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.5.1+cu121)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (0.20.1+cu121)\n",
            "Collecting qwen-vl-utils\n",
            "  Downloading qwen_vl_utils-0.0.8-py3-none-any.whl.metadata (3.6 kB)\n",
            "Collecting av\n",
            "  Downloading av-14.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.5 kB)\n",
            "Requirement already satisfied: ipython in /usr/local/lib/python3.10/dist-packages (7.34.0)\n",
            "Collecting reportlab\n",
            "  Downloading reportlab-4.2.5-py3-none-any.whl.metadata (1.5 kB)\n",
            "Collecting fpdf\n",
            "  Downloading fpdf-1.7.2.tar.gz (39 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting python-docx\n",
            "  Downloading python_docx-1.1.2-py3-none-any.whl.metadata (2.0 kB)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.10/dist-packages (11.0.0)\n",
            "Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.10/dist-packages (0.27.0)\n",
            "Collecting aiofiles<24.0,>=22.0 (from gradio)\n",
            "  Downloading aiofiles-23.2.1-py3-none-any.whl.metadata (9.7 kB)\n",
            "Requirement already satisfied: anyio<5.0,>=3.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.7.1)\n",
            "Collecting fastapi<1.0,>=0.115.2 (from gradio)\n",
            "  Downloading fastapi-0.115.6-py3-none-any.whl.metadata (27 kB)\n",
            "Collecting ffmpy (from gradio)\n",
            "  Downloading ffmpy-0.5.0-py3-none-any.whl.metadata (3.0 kB)\n",
            "Collecting gradio-client==1.5.2 (from gradio)\n",
            "  Downloading gradio_client-1.5.2-py3-none-any.whl.metadata (7.1 kB)\n",
            "Requirement already satisfied: httpx>=0.24.1 in /usr/local/lib/python3.10/dist-packages (from gradio) (0.28.1)\n",
            "Requirement already satisfied: jinja2<4.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.1.4)\n",
            "Collecting markupsafe~=2.0 (from gradio)\n",
            "  Downloading MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.0 kB)\n",
            "Requirement already satisfied: orjson~=3.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.10.12)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from gradio) (24.2)\n",
            "Requirement already satisfied: pandas<3.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (2.2.2)\n",
            "Requirement already satisfied: pydantic>=2.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (2.10.3)\n",
            "Collecting pydub (from gradio)\n",
            "  Downloading pydub-0.25.1-py2.py3-none-any.whl.metadata (1.4 kB)\n",
            "Collecting python-multipart>=0.0.18 (from gradio)\n",
            "  Downloading python_multipart-0.0.20-py3-none-any.whl.metadata (1.8 kB)\n",
            "Requirement already satisfied: pyyaml<7.0,>=5.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (6.0.2)\n",
            "Collecting ruff>=0.2.2 (from gradio)\n",
            "  Downloading ruff-0.8.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (25 kB)\n",
            "Collecting safehttpx<0.2.0,>=0.1.6 (from gradio)\n",
            "  Downloading safehttpx-0.1.6-py3-none-any.whl.metadata (4.2 kB)\n",
            "Collecting semantic-version~=2.0 (from gradio)\n",
            "  Downloading semantic_version-2.10.0-py2.py3-none-any.whl.metadata (9.7 kB)\n",
            "Collecting starlette<1.0,>=0.40.0 (from gradio)\n",
            "  Downloading starlette-0.45.1-py3-none-any.whl.metadata (6.3 kB)\n",
            "Collecting tomlkit<0.14.0,>=0.12.0 (from gradio)\n",
            "  Downloading tomlkit-0.13.2-py3-none-any.whl.metadata (2.7 kB)\n",
            "Requirement already satisfied: typer<1.0,>=0.12 in /usr/local/lib/python3.10/dist-packages (from gradio) (0.15.1)\n",
            "Requirement already satisfied: typing-extensions~=4.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (4.12.2)\n",
            "Collecting uvicorn>=0.14.0 (from gradio)\n",
            "  Downloading uvicorn-0.34.0-py3-none-any.whl.metadata (6.5 kB)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from gradio-client==1.5.2->gradio) (2024.10.0)\n",
            "Requirement already satisfied: websockets<15.0,>=10.0 in /usr/local/lib/python3.10/dist-packages (from gradio-client==1.5.2->gradio) (14.1)\n",
            "Requirement already satisfied: psutil<6,>=2 in /usr/local/lib/python3.10/dist-packages (from spaces) (5.9.5)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.16.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2024.11.6)\n",
            "Requirement already satisfied: tokenizers<0.22,>=0.21 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.21.0)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.5)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.67.1)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests) (2024.12.14)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.4.2)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch) (1.3.0)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.10/dist-packages (from ipython) (75.1.0)\n",
            "Collecting jedi>=0.16 (from ipython)\n",
            "  Downloading jedi-0.19.2-py2.py3-none-any.whl.metadata (22 kB)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from ipython) (4.4.2)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.10/dist-packages (from ipython) (0.7.5)\n",
            "Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.10/dist-packages (from ipython) (5.7.1)\n",
            "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from ipython) (3.0.48)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.10/dist-packages (from ipython) (2.18.0)\n",
            "Requirement already satisfied: backcall in /usr/local/lib/python3.10/dist-packages (from ipython) (0.2.0)\n",
            "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from ipython) (0.1.7)\n",
            "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.10/dist-packages (from ipython) (4.9.0)\n",
            "Requirement already satisfied: chardet in /usr/local/lib/python3.10/dist-packages (from reportlab) (5.2.0)\n",
            "Requirement already satisfied: lxml>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from python-docx) (5.3.0)\n",
            "Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.10/dist-packages (from anyio<5.0,>=3.0->gradio) (1.3.1)\n",
            "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<5.0,>=3.0->gradio) (1.2.2)\n",
            "Collecting starlette<1.0,>=0.40.0 (from gradio)\n",
            "  Downloading starlette-0.41.3-py3-none-any.whl.metadata (6.0 kB)\n",
            "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx>=0.24.1->gradio) (1.0.7)\n",
            "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.10/dist-packages (from httpcore==1.*->httpx>=0.24.1->gradio) (0.14.0)\n",
            "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.10/dist-packages (from jedi>=0.16->ipython) (0.8.4)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0,>=1.0->gradio) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0,>=1.0->gradio) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0,>=1.0->gradio) (2024.2)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.10/dist-packages (from pexpect>4.3->ipython) (0.7.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.10/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython) (0.2.13)\n",
            "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from pydantic>=2.0->gradio) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.27.1 in /usr/local/lib/python3.10/dist-packages (from pydantic>=2.0->gradio) (2.27.1)\n",
            "Requirement already satisfied: click>=8.0.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0,>=0.12->gradio) (8.1.7)\n",
            "Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0,>=0.12->gradio) (1.5.4)\n",
            "Requirement already satisfied: rich>=10.11.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0,>=0.12->gradio) (13.9.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas<3.0,>=1.0->gradio) (1.17.0)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=10.11.0->typer<1.0,>=0.12->gradio) (3.0.0)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer<1.0,>=0.12->gradio) (0.1.2)\n",
            "Downloading gradio-5.9.1-py3-none-any.whl (57.2 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.2/57.2 MB\u001b[0m \u001b[31m14.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading gradio_client-1.5.2-py3-none-any.whl (320 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m320.4/320.4 kB\u001b[0m \u001b[31m29.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading spaces-0.31.1-py3-none-any.whl (28 kB)\n",
            "Downloading qwen_vl_utils-0.0.8-py3-none-any.whl (5.9 kB)\n",
            "Downloading av-14.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (33.0 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m33.0/33.0 MB\u001b[0m \u001b[31m62.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading reportlab-4.2.5-py3-none-any.whl (1.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m84.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading python_docx-1.1.2-py3-none-any.whl (244 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m244.3/244.3 kB\u001b[0m \u001b[31m22.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading aiofiles-23.2.1-py3-none-any.whl (15 kB)\n",
            "Downloading fastapi-0.115.6-py3-none-any.whl (94 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m94.8/94.8 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jedi-0.19.2-py2.py3-none-any.whl (1.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m75.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (25 kB)\n",
            "Downloading python_multipart-0.0.20-py3-none-any.whl (24 kB)\n",
            "Downloading ruff-0.8.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (11.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.3/11.3 MB\u001b[0m \u001b[31m114.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading safehttpx-0.1.6-py3-none-any.whl (8.7 kB)\n",
            "Downloading semantic_version-2.10.0-py2.py3-none-any.whl (15 kB)\n",
            "Downloading starlette-0.41.3-py3-none-any.whl (73 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m73.2/73.2 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading tomlkit-0.13.2-py3-none-any.whl (37 kB)\n",
            "Downloading uvicorn-0.34.0-py3-none-any.whl (62 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.3/62.3 kB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading ffmpy-0.5.0-py3-none-any.whl (6.0 kB)\n",
            "Downloading pydub-0.25.1-py2.py3-none-any.whl (32 kB)\n",
            "Building wheels for collected packages: fpdf\n",
            "  Building wheel for fpdf (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for fpdf: filename=fpdf-1.7.2-py2.py3-none-any.whl size=40704 sha256=442a41ba3b572ac9bae1220ac5f13b34f02252630b892bab6e55af0e878115c0\n",
            "  Stored in directory: /root/.cache/pip/wheels/f9/95/ba/f418094659025eb9611f17cbcaf2334236bf39a0c3453ea455\n",
            "Successfully built fpdf\n",
            "Installing collected packages: pydub, fpdf, uvicorn, tomlkit, semantic-version, ruff, reportlab, python-multipart, python-docx, markupsafe, jedi, ffmpy, av, aiofiles, starlette, qwen-vl-utils, safehttpx, gradio-client, fastapi, gradio, spaces\n",
            "  Attempting uninstall: markupsafe\n",
            "    Found existing installation: MarkupSafe 3.0.2\n",
            "    Uninstalling MarkupSafe-3.0.2:\n",
            "      Successfully uninstalled MarkupSafe-3.0.2\n",
            "Successfully installed aiofiles-23.2.1 av-14.0.1 fastapi-0.115.6 ffmpy-0.5.0 fpdf-1.7.2 gradio-5.9.1 gradio-client-1.5.2 jedi-0.19.2 markupsafe-2.1.5 pydub-0.25.1 python-docx-1.1.2 python-multipart-0.0.20 qwen-vl-utils-0.0.8 reportlab-4.2.5 ruff-0.8.5 safehttpx-0.1.6 semantic-version-2.10.0 spaces-0.31.1 starlette-0.41.3 tomlkit-0.13.2 uvicorn-0.34.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Authenticate with Hugging Face\n",
        "from huggingface_hub import login\n",
        "\n",
        "# Log in to Hugging Face using the provided token\n",
        "hf_token = '----xxx----'\n",
        "login(hf_token)\n",
        "\n",
        "#Demo\n",
        "import gradio as gr\n",
        "import spaces\n",
        "from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer\n",
        "from qwen_vl_utils import process_vision_info\n",
        "import torch\n",
        "from PIL import Image\n",
        "import os\n",
        "import uuid\n",
        "import io\n",
        "from threading import Thread\n",
        "from reportlab.lib.pagesizes import A4\n",
        "from reportlab.lib.styles import getSampleStyleSheet\n",
        "from reportlab.lib import colors\n",
        "from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer\n",
        "from reportlab.lib.units import inch\n",
        "from reportlab.pdfbase import pdfmetrics\n",
        "from reportlab.pdfbase.ttfonts import TTFont\n",
        "import docx\n",
        "from docx.enum.text import WD_ALIGN_PARAGRAPH\n",
        "\n",
        "# Define model options\n",
        "MODEL_OPTIONS = {\n",
        "    \"OCR-KIE\": \"prithivMLmods/Qwen2-VL-OCR-2B-Instruct\",\n",
        "}\n",
        "\n",
        "# Preload models and processors into CUDA\n",
        "models = {}\n",
        "processors = {}\n",
        "for name, model_id in MODEL_OPTIONS.items():\n",
        "    print(f\"Loading {name}...\")\n",
        "    models[name] = Qwen2VLForConditionalGeneration.from_pretrained(\n",
        "        model_id,\n",
        "        trust_remote_code=True,\n",
        "        torch_dtype=torch.float16\n",
        "    ).to(\"cuda\").eval()\n",
        "    processors[name] = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
        "\n",
        "image_extensions = Image.registered_extensions()\n",
        "\n",
        "def identify_and_save_blob(blob_path):\n",
        "    \"\"\"Identifies if the blob is an image and saves it.\"\"\"\n",
        "    try:\n",
        "        with open(blob_path, 'rb') as file:\n",
        "            blob_content = file.read()\n",
        "            try:\n",
        "                Image.open(io.BytesIO(blob_content)).verify()  # Check if it's a valid image\n",
        "                extension = \".png\"  # Default to PNG for saving\n",
        "                media_type = \"image\"\n",
        "            except (IOError, SyntaxError):\n",
        "                raise ValueError(\"Unsupported media type. Please upload a valid image.\")\n",
        "\n",
        "            filename = f\"temp_{uuid.uuid4()}_media{extension}\"\n",
        "            with open(filename, \"wb\") as f:\n",
        "                f.write(blob_content)\n",
        "\n",
        "            return filename, media_type\n",
        "\n",
        "    except FileNotFoundError:\n",
        "        raise ValueError(f\"The file {blob_path} was not found.\")\n",
        "    except Exception as e:\n",
        "        raise ValueError(f\"An error occurred while processing the file: {e}\")\n",
        "\n",
        "@spaces.GPU\n",
        "def qwen_inference(model_name, media_input, text_input=None):\n",
        "    \"\"\"Handles inference for the selected model.\"\"\"\n",
        "    model = models[model_name]\n",
        "    processor = processors[model_name]\n",
        "\n",
        "    if isinstance(media_input, str):\n",
        "        media_path = media_input\n",
        "        if media_path.endswith(tuple([i for i in image_extensions.keys()])):\n",
        "            media_type = \"image\"\n",
        "        else:\n",
        "            try:\n",
        "                media_path, media_type = identify_and_save_blob(media_input)\n",
        "            except Exception as e:\n",
        "                raise ValueError(\"Unsupported media type. Please upload a valid image.\")\n",
        "\n",
        "    messages = [\n",
        "        {\n",
        "            \"role\": \"user\",\n",
        "            \"content\": [\n",
        "                {\n",
        "                    \"type\": media_type,\n",
        "                    media_type: media_path\n",
        "                },\n",
        "                {\"type\": \"text\", \"text\": text_input},\n",
        "            ],\n",
        "        }\n",
        "    ]\n",
        "\n",
        "    text = processor.apply_chat_template(\n",
        "        messages, tokenize=False, add_generation_prompt=True\n",
        "    )\n",
        "    image_inputs, _ = process_vision_info(messages)\n",
        "    inputs = processor(\n",
        "        text=[text],\n",
        "        images=image_inputs,\n",
        "        padding=True,\n",
        "        return_tensors=\"pt\",\n",
        "    ).to(\"cuda\")\n",
        "\n",
        "    streamer = TextIteratorStreamer(\n",
        "        processor.tokenizer, skip_prompt=True, skip_special_tokens=True\n",
        "    )\n",
        "    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)\n",
        "\n",
        "    thread = Thread(target=model.generate, kwargs=generation_kwargs)\n",
        "    thread.start()\n",
        "\n",
        "    buffer = \"\"\n",
        "    for new_text in streamer:\n",
        "        buffer += new_text\n",
        "        # Remove <|im_end|> or similar tokens from the output\n",
        "        buffer = buffer.replace(\"<|im_end|>\", \"\")\n",
        "        yield buffer\n",
        "\n",
        "def format_plain_text(output_text):\n",
        "    \"\"\"Formats the output text as plain text without LaTeX delimiters.\"\"\"\n",
        "    # Remove LaTeX delimiters and convert to plain text\n",
        "    plain_text = output_text.replace(\"\\\\(\", \"\").replace(\"\\\\)\", \"\").replace(\"\\\\[\", \"\").replace(\"\\\\]\", \"\")\n",
        "    return plain_text\n",
        "\n",
        "def generate_document(media_path, output_text, file_format, font_size, line_spacing, alignment, image_size):\n",
        "    \"\"\"Generates a document with the input image and plain text output.\"\"\"\n",
        "    plain_text = format_plain_text(output_text)\n",
        "    if file_format == \"pdf\":\n",
        "        return generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size)\n",
        "    elif file_format == \"docx\":\n",
        "        return generate_docx(media_path, plain_text, font_size, line_spacing, alignment, image_size)\n",
        "\n",
        "def generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size):\n",
        "    \"\"\"Generates a PDF document.\"\"\"\n",
        "    filename = f\"output_{uuid.uuid4()}.pdf\"\n",
        "    doc = SimpleDocTemplate(\n",
        "        filename,\n",
        "        pagesize=A4,\n",
        "        rightMargin=inch,\n",
        "        leftMargin=inch,\n",
        "        topMargin=inch,\n",
        "        bottomMargin=inch\n",
        "    )\n",
        "    styles = getSampleStyleSheet()\n",
        "    styles[\"Normal\"].fontSize = int(font_size)\n",
        "    styles[\"Normal\"].leading = int(font_size) * line_spacing\n",
        "    styles[\"Normal\"].alignment = {\n",
        "        \"Left\": 0,\n",
        "        \"Center\": 1,\n",
        "        \"Right\": 2,\n",
        "        \"Justified\": 4\n",
        "    }[alignment]\n",
        "\n",
        "    story = []\n",
        "\n",
        "    # Add image with size adjustment\n",
        "    image_sizes = {\n",
        "        \"Small\": (200, 200),\n",
        "        \"Medium\": (400, 400),\n",
        "        \"Large\": (600, 600)\n",
        "    }\n",
        "    img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1])\n",
        "    story.append(img)\n",
        "    story.append(Spacer(1, 12))\n",
        "\n",
        "    # Add plain text output\n",
        "    text = Paragraph(plain_text, styles[\"Normal\"])\n",
        "    story.append(text)\n",
        "\n",
        "    doc.build(story)\n",
        "    return filename\n",
        "\n",
        "def generate_docx(media_path, plain_text, font_size, line_spacing, alignment, image_size):\n",
        "    \"\"\"Generates a DOCX document.\"\"\"\n",
        "    filename = f\"output_{uuid.uuid4()}.docx\"\n",
        "    doc = docx.Document()\n",
        "\n",
        "    # Add image with size adjustment\n",
        "    image_sizes = {\n",
        "        \"Small\": docx.shared.Inches(2),\n",
        "        \"Medium\": docx.shared.Inches(4),\n",
        "        \"Large\": docx.shared.Inches(6)\n",
        "    }\n",
        "    doc.add_picture(media_path, width=image_sizes[image_size])\n",
        "    doc.add_paragraph()\n",
        "\n",
        "    # Add plain text output\n",
        "    paragraph = doc.add_paragraph()\n",
        "    paragraph.paragraph_format.line_spacing = line_spacing\n",
        "    paragraph.paragraph_format.alignment = {\n",
        "        \"Left\": WD_ALIGN_PARAGRAPH.LEFT,\n",
        "        \"Center\": WD_ALIGN_PARAGRAPH.CENTER,\n",
        "        \"Right\": WD_ALIGN_PARAGRAPH.RIGHT,\n",
        "        \"Justified\": WD_ALIGN_PARAGRAPH.JUSTIFY\n",
        "    }[alignment]\n",
        "    run = paragraph.add_run(plain_text)\n",
        "    run.font.size = docx.shared.Pt(int(font_size))\n",
        "\n",
        "    doc.save(filename)\n",
        "    return filename\n",
        "\n",
        "# CSS for output styling\n",
        "css = \"\"\"\n",
        "  #output {\n",
        "    height: 500px;\n",
        "    overflow: auto;\n",
        "    border: 1px solid #ccc;\n",
        "  }\n",
        ".submit-btn {\n",
        "    background-color: #cf3434  !important;\n",
        "    color: white !important;\n",
        "}\n",
        ".submit-btn:hover {\n",
        "    background-color: #ff2323 !important;\n",
        "}\n",
        ".download-btn {\n",
        "    background-color: #35a6d6 !important;\n",
        "    color: white !important;\n",
        "}\n",
        ".download-btn:hover {\n",
        "    background-color: #22bcff !important;\n",
        "}\n",
        "\"\"\"\n",
        "\n",
        "# Gradio app setup\n",
        "with gr.Blocks(css=css) as demo:\n",
        "    gr.Markdown(\"# Qwen2VL Models: Vision and Language Processing\")\n",
        "\n",
        "    with gr.Tab(label=\"Image Input\"):\n",
        "\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                model_choice = gr.Dropdown(\n",
        "                    label=\"Model Selection\",\n",
        "                    choices=list(MODEL_OPTIONS.keys()),\n",
        "                    value=\"OCR-KIE\"\n",
        "                )\n",
        "                input_media = gr.File(\n",
        "                    label=\"Upload Image\", type=\"filepath\"\n",
        "                )\n",
        "                text_input = gr.Textbox(label=\"Question\", placeholder=\"Ask a question about the image...\")\n",
        "                submit_btn = gr.Button(value=\"Submit\", elem_classes=\"submit-btn\")\n",
        "\n",
        "            with gr.Column():\n",
        "                output_text = gr.Textbox(label=\"Output Text\", lines=10)\n",
        "                plain_text_output = gr.Textbox(label=\"Standardized Plain Text\", lines=10)\n",
        "\n",
        "        submit_btn.click(\n",
        "            qwen_inference, [model_choice, input_media, text_input], [output_text]\n",
        "        ).then(\n",
        "            lambda output_text: format_plain_text(output_text), [output_text], [plain_text_output]\n",
        "        )\n",
        "\n",
        "        # Add examples directly usable by clicking\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                line_spacing = gr.Dropdown(\n",
        "                    choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0],\n",
        "                    value=1.5,\n",
        "                    label=\"Line Spacing\"\n",
        "                )\n",
        "                font_size = gr.Dropdown(\n",
        "                    choices=[\"8\", \"10\", \"12\", \"14\", \"16\", \"18\", \"20\", \"22\", \"24\"],\n",
        "                    value=\"18\",\n",
        "                    label=\"Font Size\"\n",
        "                )\n",
        "                alignment = gr.Dropdown(\n",
        "                    choices=[\"Left\", \"Center\", \"Right\", \"Justified\"],\n",
        "                    value=\"Justified\",\n",
        "                    label=\"Text Alignment\"\n",
        "                )\n",
        "                image_size = gr.Dropdown(\n",
        "                    choices=[\"Small\", \"Medium\", \"Large\"],\n",
        "                    value=\"Small\",\n",
        "                    label=\"Image Size\"\n",
        "                )\n",
        "                file_format = gr.Radio([\"pdf\", \"docx\"], label=\"File Format\", value=\"pdf\")\n",
        "                get_document_btn = gr.Button(value=\"Get Document\", elem_classes=\"download-btn\")\n",
        "\n",
        "        get_document_btn.click(\n",
        "            generate_document, [input_media, output_text, file_format, font_size, line_spacing, alignment, image_size], gr.File(label=\"Download Document\")\n",
        "        )\n",
        "\n",
        "demo.launch(debug=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 715
        },
        "id": "ovBSsRFhGbs2",
        "outputId": "2a5bc724-4eab-4167-9c52-b378ce3799e3"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading OCR-KIE...\n",
            "Running Gradio in a Colab notebook requires sharing enabled. Automatically setting `share=True` (you can turn this off by setting `share=False` in `launch()` explicitly).\n",
            "\n",
            "Colab notebook detected. This cell will run indefinitely so that you can see errors and logs. To turn off, set debug=False in launch().\n",
            "* Running on public URL: https://49b22c4dd53a6ec06e.gradio.live\n",
            "\n",
            "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<div><iframe src=\"https://49b22c4dd53a6ec06e.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Keyboard interruption in main thread... closing server.\n",
            "Killing tunnel 127.0.0.1:7860 <> https://49b22c4dd53a6ec06e.gradio.live\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": []
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    }
  ]
}