ahmedyasser006/arabic-legal-ocr-gemma-3-4b-qlora

pip install transformers==4.57.6 accelerate==1.8.0 peft==0.17.1 json-repair pillow

pip install -q transformers==4.57.6 vllm==0.15.0 json-repair

import base64
from io import BytesIO
from PIL import Image, ImageEnhance

def preprocess_image(image_path, max_width=1024, do_enhance=True, return_base64=False):
    image = Image.open(image_path)
    gray_image = image.convert('L')

    if gray_image.width > max_width:
        ratio = max_width / float(gray_image.width)
        new_height = int(gray_image.height * ratio)
        gray_image = gray_image.resize((max_width, new_height), Image.LANCZOS)

    if do_enhance:
        enhancer = ImageEnhance.Contrast(gray_image)
        gray_image = enhancer.enhance(1.5)

    if return_base64:
        buffered = BytesIO()
        gray_image.save(buffered, format="JPEG", optimize=True, quality=95)
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return f"data:image/jpeg;base64,{img_str}"

    return gray_image

import torch
import json_repair
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration
from peft import PeftModel

BASE_MODEL = "google/gemma-3-4b-it"
ADAPTER = "ahmedyasser006/arabic-legal-ocr-gemma-3-4b-qlora"

processor = AutoProcessor.from_pretrained(BASE_MODEL)

base_model = Gemma3ForConditionalGeneration.from_pretrained(
    BASE_MODEL, dtype="auto", device_map="auto"
)
model = PeftModel.from_pretrained(base_model, ADAPTER)
model.eval()

image = preprocess_image("document.jpg", return_base64=False)

task_1_message = """
You are a professional OCR Details Extractor.
Your rule to extract: the page markdown content in addition to the structural_elements of the document.
Extract the final output into a json format.
Do not generate any introduction or conclusion.
""".strip()


task_2_message = """
You are a professional OCR Details Extractor.
Your rule to extract the: document_classification, source, physical_properties, official_marks, signatures_authorization, routing_distribution, attachments_references, condition_notes and confidence_quality of the document.
Extract the final output into a json format.
Do not generate any introduction or conclusion.
""".strip()


messages = [
    {
        "role": "system",
        "content": [
            {
                "type": "text",
                "text": "You are a helpful assistant."
            }
        ],
    },
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": image,
            },
            {
                "type": "text",
                "text": task_1_message,
            },
        ],
    },
]


inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt",
)


# Move tensor inputs to the model device
inputs = {
    key: value.to(model.device)
    if hasattr(value, "to")
    else value
    for key, value in inputs.items()
}


input_length = inputs["input_ids"].shape[-1]


with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=False,
    )


generated_tokens = outputs[0,input_length:]


result = processor.decode(
    generated_tokens,
    skip_special_tokens=True,
)

# Robust JSON parsing, even if the model output isn't perfectly formed
json_data = json_repair.loads(result)
print(result)

vllm serve google/gemma-3-4b-it \
  --enable-lora \
  --lora-modules arabic-legal-ocr=ahmedyasser006/arabic-legal-ocr-gemma-3-4b-qlora \
  --dtype bfloat16 --gpu_memory_utilization 0.8 \
  --enable-chunked-prefill \
  --allowed-local-media-path "/workspace/"

from openai import OpenAI
import json_repair

client = OpenAI(api_key="any", base_url="http://localhost:8000/v1")

b64_image = preprocess_image("document.jpg", return_base64=True)

response = client.chat.completions.create(
    model="arabic-legal-ocr",
    messages=[{"role": "user", "content": [
        {"type": "image_url", "image_url": {"url": b64_image}},
        {"type": "text", "text": "Extract details to JSON."}
    ]}]
)

structured_output = json_repair.loads(response.choices[0].message.content)
print(structured_output)

{
  "classification": { "type": null, "category": null, "language": "Arabic" },
  "source": { "authority": null, "document_number": null, "primary_date": null },
  "content": { "subject": null, "full_text": "", "tables": [], "legal_articles": [] },
  "quality": { "confidence": null, "manual_review": false }
}

arabic-legal-ocr-gemma-3-4b-qlora/
├── README.md
├── adapter_config.json
├── adapter_model.safetensors     ← final adapter, use for inference
├── tokenizer.model / tokenizer_config.json / ...
├── train_results.json / eval_results.json / trainer_state.json
└── last-checkpoint/               ← training-resume checkpoint only
    ├── adapter_model.safetensors
    ├── optimizer.pt / scheduler.pt / scaler.pt / rng_state*.pth
    └── trainer_state.json