lukaszwysocki/lightonocr2-cpu repository overview

FROM python:3.12-slim

# Instalacja zależności systemowych
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Poprawiona instalacja:
# 1. Najpierw instalujemy torch w wersji CPU
# 2. Potem resztę z domyślnego PyPI
RUN pip install --no-cache-dir torch --extra-index-url https://download.pytorch.org/whl/cpu && \
    pip install --no-cache-dir \
    git+https://github.com/huggingface/transformers \
    accelerate \
    pypdfium2 \
    pillow \
    sentencepiece

CMD ["python3"]

import torch
from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
import pypdfium2 as pdfium
from PIL import Image
import os

MODEL_ID = "lightonai/LightOnOCR-2-1B"
PDF_PATH = "stary_scan.pdf"
PAGE_INDEX = 1 

def run_ocr_cpu():
    print("Inicjalizacja na CPU (może to zająć chwilę)...")
    device = "cpu"
    
    # Ładowanie modelu w float32 (bezpieczniejsze i dokładniejsze na CPU)
    model = LightOnOcrForConditionalGeneration.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float32, 
        device_map={"": "cpu"} 
    )
    processor = LightOnOcrProcessor.from_pretrained(MODEL_ID)

    # Renderowanie PDF
    pdf = pdfium.PdfDocument(PDF_PATH)
    try:
        page = pdf[PAGE_INDEX]
        bitmap = page.render(scale=2.0) # Zmniejszono nieco scale dla szybkości na CPU
        pil_image = bitmap.to_pil()
    finally:
        pdf.close()

    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Convert this document page to markdown."}
            ]
        }
    ]
    
    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
    # Na CPU nie używamy bfloat16
    inputs = processor(images=pil_image, text=prompt, return_tensors="pt").to(device)

    print("Generowanie OCR (na CPU będzie to trwać od 30s do kilku minut)...")
    with torch.no_grad():
        output_ids = model.generate(
            **inputs, 
            max_new_tokens=4096, # Umiarkowany limit dla CPU
            do_sample=False,
            use_cache=True
        )
    
    generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
    result_text = processor.decode(generated_ids, skip_special_tokens=True)

    with open("wynik_cpu_strona_2.md", "w", encoding="utf-8") as f:
        f.write(result_text)
    
    print("Gotowe! Wynik zapisano w wynik_cpu_strona_2.md")

if __name__ == "__main__":
    run_ocr_cpu()