From da2f4779c18e0c26905720839ba5cd9e4a8e54b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Pin=C3=A7on?= Date: Tue, 24 Jun 2025 10:15:52 +0200 Subject: [PATCH 1/2] init ci --- .github/workflows/release.yml | 40 +++++++++ Dockerfile | 6 ++ backend/config.py | 2 +- backend/page_treatment/analyzer.py | 135 +++++++++++++---------------- backend/requirements.txt | 7 ++ docker-compose.yml | 7 ++ 6 files changed, 119 insertions(+), 78 deletions(-) create mode 100644 .github/workflows/release.yml create mode 100644 Dockerfile create mode 100644 backend/requirements.txt create mode 100644 docker-compose.yml diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..61df7d8 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,40 @@ +name: Build and Push Docker Image + +on: + push: + branches: + - main + - docker + +jobs: + build-and-push: + runs-on: ubuntu-latest + + permissions: + contents: read + packages: write + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Set image name + id: image_name + run: | + IMAGE_NAME=ghcr.io/${{ github.repository }}:${{ github.sha }} + echo "IMAGE_NAME=${IMAGE_NAME,,}" >> "$GITHUB_OUTPUT" + + - name: Build Docker image + run: | + docker build -t ${{ steps.image_name.outputs.IMAGE_NAME }} . + + - name: Push Docker image + run: | + docker push ${{ steps.image_name.outputs.IMAGE_NAME }} diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..a0032b9 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,6 @@ +FROM python:3 +COPY . . +WORKDIR backend +RUN pip install -r requirements.txt +WORKDIR .. +ENTRYPOINT run_app.sh \ No newline at end of file diff --git a/backend/config.py b/backend/config.py index c14bfaa..4173565 100644 --- a/backend/config.py +++ b/backend/config.py @@ -36,7 +36,7 @@ CLEANUP_AGE_HOURS = 24 CLEANUP_INTERVAL = 3600 # 1 hour # Model settings -MODEL_PATH = "ds4sd/SmolDocling-256M-preview-mlx-bf16" +MODEL_PATH = "ds4sd/SmolDocling-256M-preview" MAX_TOKENS = 4096 # Zone colors for visualization diff --git a/backend/page_treatment/analyzer.py b/backend/page_treatment/analyzer.py index d6bdb4b..36762d3 100644 --- a/backend/page_treatment/analyzer.py +++ b/backend/page_treatment/analyzer.py @@ -2,36 +2,37 @@ # /// script # requires-python = ">=3.12" # dependencies = [ -# "docling-core", -# "mlx-vlm", +# "transformers>=4.50", +# "torch", # "pillow", # "requests", # "argparse", # "pdf2image", +# "docling_core", # ] # /// import argparse import os import tempfile -import re from pathlib import Path from urllib.parse import urlparse import requests from PIL import Image from pdf2image import convert_from_bytes -from docling_core.types.doc import ImageRefMode -from docling_core.types.doc.document import DocTagsDocument, DoclingDocument +import torch +from transformers import AutoProcessor, AutoModelForVision2Seq +from docling_core.types.doc import DoclingDocument +from docling_core.types.doc.document import DocTagsDocument # Add parent directory to path for imports import sys sys.path.append(str(Path(__file__).parent.parent.parent)) -from backend.utils import ensure_results_folder, load_pdf_page, get_project_root +from backend.utils import ensure_results_folder, load_pdf_page from backend.config import MODEL_PATH, MAX_TOKENS, DEFAULT_DPI def parse_arguments(): - """Parse command line arguments.""" results_dir = ensure_results_folder() parser = argparse.ArgumentParser(description='Convert an image or PDF to docling format') @@ -52,7 +53,6 @@ def parse_arguments(): return parser.parse_args() def load_image(image_path, page_num=1, dpi=DEFAULT_DPI): - """Load image from URL, local image file, or PDF.""" if urlparse(image_path).scheme in ['http', 'https']: response = requests.get(image_path, stream=True, timeout=10) response.raise_for_status() @@ -62,104 +62,85 @@ def load_image(image_path, page_num=1, dpi=DEFAULT_DPI): pdf_images = convert_from_bytes(response.content, dpi=dpi, first_page=page_num, last_page=page_num) if not pdf_images: raise Exception(f"Could not extract page {page_num} from PDF") - return pdf_images[0] + return pdf_images[0].convert("RGB") else: - return Image.open(response.raw) + return Image.open(response.raw).convert("RGB") else: image_path = Path(image_path) if not image_path.exists(): raise FileNotFoundError(f"File not found: {image_path}") if image_path.suffix.lower() == '.pdf': - return load_pdf_page(str(image_path), page_num, dpi) + return load_pdf_page(str(image_path), page_num, dpi).convert("RGB") else: - return Image.open(image_path) - -def process_page(model, processor, config, args, pil_image, page_num=1): - """Process a single page from a PDF or image file.""" - from mlx_vlm.prompt_utils import apply_chat_template - from mlx_vlm.utils import stream_generate + return Image.open(image_path).convert("RGB") +def process_page(model, processor, args, pil_image, page_num=1): results_dir = ensure_results_folder() - # For web interface, always use output.doctags.txt - # For command line with specific pages, use page-specific names if args.start_page == args.end_page and args.start_page == page_num: - # Single page processing - output_path = results_dir / "output.html" doctags_path = results_dir / "output.doctags.txt" + output_path = results_dir / "output.html" else: - # Multi-page processing - output_path = results_dir / f"output_page{page_num}.html" doctags_path = results_dir / f"output_page{page_num}.doctags.txt" + output_path = results_dir / f"output_page{page_num}.html" print(f"Processing page {page_num}") - # Save image temporarily - with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img_file: - temp_img_path = temp_img_file.name - pil_image.save(temp_img_path, format='PNG') + # Préparer les messages + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": args.prompt} + ] + } + ] - try: - # Apply chat template and generate - formatted_prompt = apply_chat_template(processor, config, args.prompt, num_images=1) + prompt = processor.apply_chat_template(messages, add_generation_prompt=True) - print(f"Generating DocTags for page {page_num}: \n\n") - output = "" - for token in stream_generate( - model, processor, formatted_prompt, [temp_img_path], max_tokens=MAX_TOKENS, verbose=False - ): - output += token.text - print(token.text, end="") - if "" in token.text: - break - print("\n\n") + device = next(model.parameters()).device + inputs = processor(text=prompt, images=[pil_image], return_tensors="pt").to(device) - finally: - # Clean up temporary file - if os.path.exists(temp_img_path): - os.unlink(temp_img_path) + # Génération + generated_ids = model.generate(**inputs, max_new_tokens=MAX_TOKENS) + prompt_length = inputs.input_ids.shape[1] + trimmed_generated_ids = generated_ids[:, prompt_length:] - # Save DocTags output - with open(doctags_path, 'w', encoding='utf-8') as f: - f.write(output) - print(f"Raw DocTags saved to: {doctags_path}") + doctags = processor.batch_decode(trimmed_generated_ids, skip_special_tokens=False)[0].lstrip() + with open(doctags_path, "w", encoding="utf-8") as f: + f.write(doctags) + print(f"DocTags saved to {doctags_path}") + + doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [pil_image]) + doc = DoclingDocument.load_from_doctags(doctags_doc, document_name=f"Page {page_num}") + html = doc.export_to_html() + + with open(output_path, "w", encoding="utf-8") as f: + f.write(html) + print(f"HTML exported to {output_path}") return output_path def main(): args = parse_arguments() + print("Loading model and processor...") - # Load the model - print("Loading model...") - try: - from mlx_vlm import load - from mlx_vlm.utils import load_config + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = AutoModelForVision2Seq.from_pretrained( + MODEL_PATH, + torch_dtype=torch.bfloat16, + _attn_implementation="flash_attention_2" if device.type == "cuda" else "eager" + ).to(device) + processor = AutoProcessor.from_pretrained(MODEL_PATH) - model, processor = load(MODEL_PATH) - config = load_config(MODEL_PATH) - except Exception as e: - print(f"Error loading model: {e}") - return + start_page = args.start_page + end_page = args.end_page or args.page - # Process the image/PDF - try: - # Handle single page or range - start_page = args.start_page - end_page = args.end_page or args.page - - for page_num in range(start_page, end_page + 1): - print(f"\nProcessing page {page_num}...") - - pil_image = load_image(args.image, page_num=page_num, dpi=args.dpi) - print(f"Page {page_num} loaded: {pil_image.size}") - - process_page(model, processor, config, args, pil_image, page_num) - - except Exception as e: - print(f"Error processing: {e}") - import traceback - traceback.print_exc() + for page_num in range(start_page, end_page + 1): + pil_image = load_image(args.image, page_num=page_num, dpi=args.dpi) + process_page(model, processor, args, pil_image, page_num) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/backend/requirements.txt b/backend/requirements.txt new file mode 100644 index 0000000..32d1b75 --- /dev/null +++ b/backend/requirements.txt @@ -0,0 +1,7 @@ +transformers +accelerate +torch +torchvision +pdf2image +pillow +requests diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..f53f1e5 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,7 @@ +services: + analyser: + image: ddd + ports: + - "8080:5000" + volumes: + - ./input:/ From 10900c346ea5a5307a998165128b429d6e186831 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Pin=C3=A7on?= Date: Tue, 24 Jun 2025 11:54:34 +0200 Subject: [PATCH 2/2] init ci --- .github/workflows/release.yml | 4 ++++ Dockerfile | 11 ++++++----- backend/requirements.txt | 1 + 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 61df7d8..958faf8 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -38,3 +38,7 @@ jobs: - name: Push Docker image run: | docker push ${{ steps.image_name.outputs.IMAGE_NAME }} + - name: Tag image as latest + run: | + docker tag ${{ steps.image_name.outputs.IMAGE_NAME }} ghcr.io/${{ github.repository }}:latest + docker push ghcr.io/${{ github.repository }}:latest \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index a0032b9..12c4c58 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,6 +1,7 @@ -FROM python:3 +FROM python:3.12 +WORKDIR /app +COPY backend/requirements.txt ./backend/requirements.txt +RUN pip install --no-cache-dir -r backend/requirements.txt COPY . . -WORKDIR backend -RUN pip install -r requirements.txt -WORKDIR .. -ENTRYPOINT run_app.sh \ No newline at end of file +RUN chmod +x run_app.sh +ENTRYPOINT ["./run_app.sh"] diff --git a/backend/requirements.txt b/backend/requirements.txt index 32d1b75..359de26 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -5,3 +5,4 @@ torchvision pdf2image pillow requests +flask \ No newline at end of file