初始化项目，由ModelHub XC社区提供模型
Model: prithivMLmods/Gliese-OCR-7B-Post1.0 Source: Original Platform
This commit is contained in:
ModelHub XC
@@ -0,0 +1,401 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "DgpubXociwNK"
|
||||
},
|
||||
"source": [
|
||||
"## **Gliese-OCR-7B-Post1.0(4-bit)**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Nb3wNhothvX7"
|
||||
},
|
||||
"source": [
|
||||
"The Gliese-OCR-7B-Post1.0 model is a fine-tuned version of Camel-Doc-OCR-062825, optimized for Document Retrieval, Content Extraction, and Analysis Recognition. Built on top of the Qwen2.5-VL architecture, this model enhances document comprehension capabilities with focused training on the Opendoc2-Analysis-Recognition dataset for superior document analysis and information extraction tasks.\n",
|
||||
"\n",
|
||||
" > This model shows significant improvements in LaTeX rendering and Markdown rendering for OCR tasks.\n",
|
||||
"\n",
|
||||
"| Image1 | Image2 |\n",
|
||||
"|--------|--------|\n",
|
||||
"|  |  |\n",
|
||||
"\n",
|
||||
"*multimodal model & notebook by: [prithivMLmods](https://huggingface.co/prithivMLmods)*"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Mk560Wx0j6PY"
|
||||
},
|
||||
"source": [
|
||||
"### **Install packages**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "qTD_dNliNS5T"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/huggingface/transformers.git \\\n",
|
||||
" git+https://github.com/huggingface/accelerate.git \\\n",
|
||||
" git+https://github.com/huggingface/peft.git \\\n",
|
||||
" transformers-stream-generator huggingface_hub albumentations \\\n",
|
||||
" pyvips-binary qwen-vl-utils sentencepiece opencv-python docling-core \\\n",
|
||||
" python-docx torchvision safetensors matplotlib num2words \\\n",
|
||||
"\n",
|
||||
"!pip install xformers requests pymupdf hf_xet spaces pyvips pillow gradio \\\n",
|
||||
" einops torch fpdf timm av decord bitsandbytes reportlab\n",
|
||||
"#Hold tight, this will take around 1-2 minutes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "uiBblyf-kLmf"
|
||||
},
|
||||
"source": [
|
||||
"### **Run Demo App**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "pgz93DfvNMfb"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import spaces\n",
|
||||
"import json\n",
|
||||
"import math\n",
|
||||
"import os\n",
|
||||
"import traceback\n",
|
||||
"from io import BytesIO\n",
|
||||
"from typing import Any, Dict, List, Optional, Tuple\n",
|
||||
"import re\n",
|
||||
"import time\n",
|
||||
"from threading import Thread\n",
|
||||
"from io import BytesIO\n",
|
||||
"import uuid\n",
|
||||
"import tempfile\n",
|
||||
"\n",
|
||||
"import gradio as gr\n",
|
||||
"import requests\n",
|
||||
"import torch\n",
|
||||
"from PIL import Image\n",
|
||||
"import fitz\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"# --- New Model Imports ---\n",
|
||||
"from transformers import (\n",
|
||||
" Qwen2_5_VLForConditionalGeneration,\n",
|
||||
" AutoProcessor,\n",
|
||||
" TextIteratorStreamer,\n",
|
||||
" BitsAndBytesConfig,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"from reportlab.lib.pagesizes import A4\n",
|
||||
"from reportlab.lib.styles import getSampleStyleSheet\n",
|
||||
"from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer\n",
|
||||
"from reportlab.lib.units import inch\n",
|
||||
"\n",
|
||||
"# --- Constants and Model Setup ---\n",
|
||||
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
|
||||
"\n",
|
||||
"print(\"CUDA_VISIBLE_DEVICES=\", os.environ.get(\"CUDA_VISIBLE_DEVICES\"))\n",
|
||||
"print(\"torch.__version__ =\", torch.__version__)\n",
|
||||
"print(\"torch.version.cuda =\", torch.version.cuda)\n",
|
||||
"print(\"cuda available:\", torch.cuda.is_available())\n",
|
||||
"print(\"cuda device count:\", torch.cuda.device_count())\n",
|
||||
"if torch.cuda.is_available():\n",
|
||||
" print(\"current device:\", torch.cuda.current_device())\n",
|
||||
" print(\"device name:\", torch.cuda.get_device_name(torch.cuda.current_device()))\n",
|
||||
"\n",
|
||||
"print(\"Using device:\", device)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# --- Model Loading (Updated for Qwen2.5-VL) ---\n",
|
||||
"\n",
|
||||
"# Define model options\n",
|
||||
"MODEL_OPTIONS = {\n",
|
||||
" \"Gliese-OCR-7B-Post1.0\": \"prithivMLmods/Gliese-OCR-7B-Post1.0\",\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"# Define 4-bit quantization configuration\n",
|
||||
"# This config will load the model in 4-bit to save VRAM.\n",
|
||||
"quantization_config = BitsAndBytesConfig(\n",
|
||||
" load_in_4bit=True,\n",
|
||||
" bnb_4bit_compute_dtype=torch.float16,\n",
|
||||
" bnb_4bit_quant_type=\"nf4\",\n",
|
||||
" bnb_4bit_use_double_quant=True,\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}🤗. This will use 4-bit quantization to save VRAM.\")\n",
|
||||
" models[name] = Qwen2_5_VLForConditionalGeneration.from_pretrained(\n",
|
||||
" model_id,\n",
|
||||
" trust_remote_code=True,\n",
|
||||
" quantization_config=quantization_config,\n",
|
||||
" device_map=\"auto\"\n",
|
||||
" )\n",
|
||||
" processors[name] = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
|
||||
"print(\"Model loaded successfully.\")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# --- PDF Generation and Preview Utility Function (Unchanged) ---\n",
|
||||
"def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str):\n",
|
||||
" \"\"\"\n",
|
||||
" Generates a PDF, saves it, and then creates image previews of its pages.\n",
|
||||
" Returns the path to the PDF and a list of paths to the preview images.\n",
|
||||
" \"\"\"\n",
|
||||
" if image is None or not text_content or not text_content.strip():\n",
|
||||
" raise gr.Error(\"Cannot generate PDF. Image or text content is missing.\")\n",
|
||||
"\n",
|
||||
" # --- 1. Generate the PDF ---\n",
|
||||
" temp_dir = tempfile.gettempdir()\n",
|
||||
" pdf_filename = os.path.join(temp_dir, f\"output_{uuid.uuid4()}.pdf\")\n",
|
||||
" doc = SimpleDocTemplate(\n",
|
||||
" pdf_filename,\n",
|
||||
" pagesize=A4,\n",
|
||||
" rightMargin=inch, leftMargin=inch,\n",
|
||||
" topMargin=inch, bottomMargin=inch\n",
|
||||
" )\n",
|
||||
" styles = getSampleStyleSheet()\n",
|
||||
" style_normal = styles[\"Normal\"]\n",
|
||||
" style_normal.fontSize = int(font_size)\n",
|
||||
" style_normal.leading = int(font_size) * line_spacing\n",
|
||||
" style_normal.alignment = {\"Left\": 0, \"Center\": 1, \"Right\": 2, \"Justified\": 4}[alignment]\n",
|
||||
"\n",
|
||||
" story = []\n",
|
||||
"\n",
|
||||
" img_buffer = BytesIO()\n",
|
||||
" image.save(img_buffer, format='PNG')\n",
|
||||
" img_buffer.seek(0)\n",
|
||||
"\n",
|
||||
" page_width, _ = A4\n",
|
||||
" available_width = page_width - 2 * inch\n",
|
||||
" image_widths = {\n",
|
||||
" \"Small\": available_width * 0.3,\n",
|
||||
" \"Medium\": available_width * 0.6,\n",
|
||||
" \"Large\": available_width * 0.9,\n",
|
||||
" }\n",
|
||||
" img_width = image_widths[image_size]\n",
|
||||
" # Create a ReportLab Image object, handling potential transparency\n",
|
||||
" img = RLImage(img_buffer, width=img_width, height=image.height * (img_width / image.width))\n",
|
||||
" story.append(img)\n",
|
||||
" story.append(Spacer(1, 12))\n",
|
||||
"\n",
|
||||
" # Clean the text for PDF generation\n",
|
||||
" cleaned_text = re.sub(r'#+\\s*', '', text_content).replace(\"*\", \"\")\n",
|
||||
" text_paragraphs = cleaned_text.split('\\n')\n",
|
||||
"\n",
|
||||
" for para in text_paragraphs:\n",
|
||||
" if para.strip():\n",
|
||||
" story.append(Paragraph(para, style_normal))\n",
|
||||
"\n",
|
||||
" doc.build(story)\n",
|
||||
"\n",
|
||||
" # --- 2. Render PDF pages as images for preview ---\n",
|
||||
" preview_images = []\n",
|
||||
" try:\n",
|
||||
" pdf_doc = fitz.open(pdf_filename)\n",
|
||||
" for page_num in range(len(pdf_doc)):\n",
|
||||
" page = pdf_doc.load_page(page_num)\n",
|
||||
" pix = page.get_pixmap(dpi=150)\n",
|
||||
" preview_img_path = os.path.join(temp_dir, f\"preview_{uuid.uuid4()}_p{page_num}.png\")\n",
|
||||
" pix.save(preview_img_path)\n",
|
||||
" preview_images.append(preview_img_path)\n",
|
||||
" pdf_doc.close()\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"Error generating PDF preview: {e}\")\n",
|
||||
"\n",
|
||||
" return pdf_filename, preview_images\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# --- Core Application Logic (Updated for Qwen2.5-VL with Streaming) ---\n",
|
||||
"@spaces.GPU\n",
|
||||
"def process_document(\n",
|
||||
" image: Image.Image,\n",
|
||||
" prompt_input: str,\n",
|
||||
" max_new_tokens: int,\n",
|
||||
" temperature: float,\n",
|
||||
" top_p: float,\n",
|
||||
" top_k: int,\n",
|
||||
" repetition_penalty: float\n",
|
||||
"):\n",
|
||||
" \"\"\"\n",
|
||||
" Main function that handles model inference for the Qwen model with streaming.\n",
|
||||
" This function is a generator, yielding text as it is generated.\n",
|
||||
" \"\"\"\n",
|
||||
" if image is None:\n",
|
||||
" yield \"Please upload an image.\", \"Please upload an image.\"\n",
|
||||
" return\n",
|
||||
" if not prompt_input or not prompt_input.strip():\n",
|
||||
" yield \"Please enter a prompt.\", \"Please enter a prompt.\"\n",
|
||||
" return\n",
|
||||
"\n",
|
||||
" model_name = \"Gliese-OCR-7B-Post1.0\"\n",
|
||||
" model = models[model_name]\n",
|
||||
" processor = processors[model_name]\n",
|
||||
"\n",
|
||||
" messages = [\n",
|
||||
" {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": [\n",
|
||||
" {\"type\": \"image\", \"image\": image},\n",
|
||||
" {\"type\": \"text\", \"text\": prompt_input},\n",
|
||||
" ],\n",
|
||||
" }\n",
|
||||
" ]\n",
|
||||
"\n",
|
||||
" text = processor.apply_chat_template(\n",
|
||||
" messages, tokenize=False, add_generation_prompt=True\n",
|
||||
" )\n",
|
||||
" inputs = processor(\n",
|
||||
" text=[text],\n",
|
||||
" images=[image],\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",
|
||||
"\n",
|
||||
" generation_kwargs = dict(\n",
|
||||
" inputs,\n",
|
||||
" streamer=streamer,\n",
|
||||
" max_new_tokens=max_new_tokens,\n",
|
||||
" temperature=temperature,\n",
|
||||
" top_p=top_p,\n",
|
||||
" top_k=top_k,\n",
|
||||
" repetition_penalty=repetition_penalty,\n",
|
||||
" do_sample=True if temperature > 0 else False,\n",
|
||||
" )\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 special tokens from the output stream\n",
|
||||
" clean_buffer = buffer.replace(\"<|im_end|>\", \"\").replace(\"<|endoftext|>\", \"\")\n",
|
||||
" yield clean_buffer, clean_buffer\n",
|
||||
"\n",
|
||||
"# --- Gradio UI Definition (Updated Title, otherwise unchanged) ---\n",
|
||||
"def create_gradio_interface():\n",
|
||||
" \"\"\"Builds and returns the Gradio web interface.\"\"\"\n",
|
||||
" css = \"\"\"\n",
|
||||
" .main-container { max-width: 1400px; margin: 0 auto; }\n",
|
||||
" .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}\n",
|
||||
" .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }\n",
|
||||
" #gallery { min-height: 400px; }\n",
|
||||
" \"\"\"\n",
|
||||
" with gr.Blocks(theme=\"bethecloud/storj_theme\", css=css) as demo:\n",
|
||||
" gr.HTML(f\"\"\"\n",
|
||||
" <div class=\"title\" style=\"text-align: center\">\n",
|
||||
" <h1>Gliese-OCR-7B-Post1.0 📄</h1>\n",
|
||||
" <p style=\"font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;\">\n",
|
||||
" Image Content Extraction and Markdown Rendering </b>\n",
|
||||
" </p>\n",
|
||||
" </div>\n",
|
||||
" \"\"\")\n",
|
||||
"\n",
|
||||
" with gr.Row():\n",
|
||||
" # Left Column (Inputs)\n",
|
||||
" with gr.Column(scale=1):\n",
|
||||
" prompt_input = gr.Textbox(label=\"Query Input\", placeholder=\"✦︎ Enter the prompt.\", value=\"Precisely OCR the Image.\")\n",
|
||||
" image_input = gr.Image(label=\"Upload Image\", type=\"pil\", sources=['upload'])\n",
|
||||
"\n",
|
||||
" with gr.Accordion(\"Advanced Settings\", open=False):\n",
|
||||
" max_new_tokens = gr.Slider(minimum=64, maximum=2048, value=1024, step=32, label=\"Max New Tokens\")\n",
|
||||
" temperature = gr.Slider(label=\"Temperature\", minimum=0.1, maximum=2.0, step=0.1, value=0.7)\n",
|
||||
" top_p = gr.Slider(label=\"Top-p (nucleus sampling)\", minimum=0.05, maximum=1.0, step=0.05, value=0.9)\n",
|
||||
" top_k = gr.Slider(label=\"Top-k\", minimum=1, maximum=100, step=1, value=50)\n",
|
||||
" repetition_penalty = gr.Slider(label=\"Repetition penalty\", minimum=1.0, maximum=2.0, step=0.05, value=1.1)\n",
|
||||
"\n",
|
||||
" with gr.Accordion(\"PDF Export Settings\", open=False):\n",
|
||||
" font_size = gr.Dropdown(choices=[\"8\", \"10\", \"12\", \"14\", \"16\", \"18\"], value=\"12\", label=\"Font Size\")\n",
|
||||
" line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label=\"Line Spacing\")\n",
|
||||
" alignment = gr.Dropdown(choices=[\"Left\", \"Center\", \"Right\", \"Justified\"], value=\"Justified\", label=\"Text Alignment\")\n",
|
||||
" image_size = gr.Dropdown(choices=[\"Small\", \"Medium\", \"Large\"], value=\"Medium\", label=\"Image Size in PDF\")\n",
|
||||
"\n",
|
||||
" process_btn = gr.Button(\"🚀 Process Image\", variant=\"primary\", elem_classes=[\"process-button\"], size=\"lg\")\n",
|
||||
" clear_btn = gr.Button(\"🗑️ Clear All\", variant=\"secondary\")\n",
|
||||
"\n",
|
||||
" # Right Column (Outputs)\n",
|
||||
" with gr.Column(scale=2):\n",
|
||||
" with gr.Tabs() as tabs:\n",
|
||||
" with gr.Tab(\"📝 Extracted Content\"):\n",
|
||||
" raw_output = gr.Textbox(label=\"Model Output\", interactive=False, lines=15, show_copy_button=True)\n",
|
||||
"\n",
|
||||
" gr.Markdown(\"[prithivMLmods🤗](https://huggingface.co/prithivMLmods)\")\n",
|
||||
"\n",
|
||||
" with gr.Tab(\"📰 Markdown Preview\"):\n",
|
||||
" with gr.Accordion(\"(Result.md)\", open=True):\n",
|
||||
" markdown_output = gr.Markdown()\n",
|
||||
"\n",
|
||||
" with gr.Tab(\"📋 PDF Preview\"):\n",
|
||||
" generate_pdf_btn = gr.Button(\"📄 Generate PDF & Render\", variant=\"primary\")\n",
|
||||
" pdf_output_file = gr.File(label=\"Download Generated PDF\", interactive=False)\n",
|
||||
" pdf_preview_gallery = gr.Gallery(label=\"PDF Page Preview\", show_label=True, elem_id=\"gallery\", columns=2, object_fit=\"contain\", height=\"auto\")\n",
|
||||
"\n",
|
||||
" # Event Handlers\n",
|
||||
" def clear_all_outputs():\n",
|
||||
" return None, \"\", \"Model output will appear here.\", \"\", None, None\n",
|
||||
"\n",
|
||||
" # The .click() event will now stream the output from the generator function\n",
|
||||
" process_btn.click(\n",
|
||||
" fn=process_document,\n",
|
||||
" inputs=[image_input, prompt_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty],\n",
|
||||
" outputs=[raw_output, markdown_output]\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" generate_pdf_btn.click(\n",
|
||||
" fn=generate_and_preview_pdf,\n",
|
||||
" inputs=[image_input, raw_output, font_size, line_spacing, alignment, image_size],\n",
|
||||
" outputs=[pdf_output_file, pdf_preview_gallery]\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" clear_btn.click(\n",
|
||||
" clear_all_outputs,\n",
|
||||
" outputs=[image_input, prompt_input, raw_output, markdown_output, pdf_output_file, pdf_preview_gallery]\n",
|
||||
" )\n",
|
||||
" return demo\n",
|
||||
"\n",
|
||||
"if __name__ == \"__main__\":\n",
|
||||
" demo = create_gradio_interface()\n",
|
||||
" # Use queue() for better handling of multiple users and streaming\n",
|
||||
" demo.queue(max_size=20).launch(share=True, show_error=True)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
Reference in New Issue
Block a user