import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from peft import PeftModel
from PIL import Image
base_model_id = "Qwen/Qwen3.5-4B"
model = AutoModelForCausalLM.from_pretrained(
base_model_id,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
)
processor = AutoProcessor.from_pretrained(base_model_id, trust_remote_code=True)
model = PeftModel.from_pretrained(model, "Yesianrohn")
model = model.merge_and_unload()
image = Image.open("document.png").convert("RGB")
system_prompt = (
"You are an expert document parser. Given an image of a document page, "
"reconstruct its source as a single complete, self-contained HTML5 "
"document. Faithfully preserve the original layout, typography, tables, "
"formulas, and visual hierarchy using inline CSS where appropriate. "
"Output only the HTML source, with no explanations, no markdown fences, "
"and no extra prose."
)
user_prompt = (
"Convert this document page into a complete HTML document. "
"Preserve the layout, headings, tables, and formulas exactly as shown. "
"Return only the HTML source."
)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": [
{"type": "image", "image": image},
{"type": "text", "text": user_prompt},
]},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=[text], images=[image], return_tensors="pt").to(model.device)
output_ids = model.generate(**inputs, max_new_tokens=10240, temperature=0.0, do_sample=False)
output_text = processor.batch_decode(output_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
print(output_text)