import gc
import json
import torch
from unsloth import FastVisionModel
from qwen_vl_utils import process_vision_info
from PIL import Image
model, tokenizer = FastVisionModel.from_pretrained(
"a-mo-yehia/contract_OCR_model",
load_in_4bit=False,
)
FastVisionModel.for_inference(model)
img1 = Image.open("page1.jpg").convert("RGB")
img2 = Image.open("page2.jpg").convert("RGB")
SYSTEM_PROMPT = """\
أنت نظام استخراج بيانات من عقود تأسيس شركات مصرية مسحوبة ضوئياً.
الوثيقة أمامك هي صفحتان من عقد تأسيس شركة مصرية مكتوبة بالعربية.
─── بنية المستند ───
- يبدأ العقد بديباجة (تمهيد) تذكر القانون المنظِّم للشركة.
- تأتي بعدها مواد مرقّمة (مادة ١، مادة ٢، ...) تحتوي على كل التفاصيل.
- اسم الشركة ونوعها: عادةً في المواد الأولى (١–٣).
- غرض الشركة ونشاطها: في مادة الغرض (المادة ٣ أو ٤ غالباً).
- العنوان: في المادة التي تذكر "المركز الرئيسي".
- رأس المال وتوزيعه بين الشركاء: في مادة واحدة تحتوي عادةً على جدول.
- مدة الشركة: في المادة التي تذكر السنوات أو الأجل.
- رقم القانون: في التمهيد أو أوائل المواد، صيغته "القانون رقم XXX لسنة XXXX".
─── قواعد الاستخراج ───
- انسخ كل نص عربي حرفياً كما يظهر في الصورة — لا تصحّح ولا تُعيد صياغة.
- إذا كانت مادة ممتدة وانتهت الصفحة قبل اكتمالها: "partial": true.
- إذا لم يوجد حقل في الصفحتين: اتركه سلسلة فارغة "".
- إذا لم تكن هناك مواد في الصفحتين: أرجع "articles": [].
- إذا احتوت الوثيقة على أي جداول، **يجب** استخراجها كاملة وتنسيقها حصرياً بصيغة Markdown (MD).
- يمنع منعاً باتاً تجاهل أي بيانات مجدولة أو دمجها كنص عادي.
- أرجع JSON صحيح فقط — بدون markdown، بدون شرح، بدون مفاتيح إضافية.\
"""
def run_inference(img1: Image.Image, img2: Image.Image, file_id: str, page_start: int) -> str:
gc.collect()
torch.cuda.empty_cache()
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": [
{"type": "image", "image": img1},
{"type": "image", "image": img2},
{"type": "text", "text": (
f"استخرج جميع البيانات من الصفحتين "
f"من عقد {file_id} وأرجعها كـ JSON كامل."
)},
]},
]
input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
image_inputs, _ = process_vision_info(messages)
inputs = tokenizer(
image_inputs,
input_text,
add_special_tokens=False,
return_tensors="pt",
).to("cuda")
with torch.no_grad():
out = model.generate(
**inputs,
max_new_tokens=4096,
use_cache=True,
do_sample=False,
repetition_penalty=1.1,
)
generated = tokenizer.decode(
out[0, inputs["input_ids"].shape[-1]:],
skip_special_tokens=True,
).strip()
return generated
raw_output = run_inference(img1, img2, FILE_ID, PAGE_START)
print("\n=== Raw Output ===")
print(raw_output)