LinBMS410/receipt_detector

from peft import PeftModel
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info
import torch

processor = AutoProcessor.from_pretrained("LinBMS410/receipt_detector")
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-VL-3B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model = PeftModel.from_pretrained(model, "LinBMS410/receipt_detector")
model.eval()

# 推論
def predict(image):
    messages = [
        {"role": "system", "content": [{"type": "text", "text": "你是專業的收據資訊抽取助手。請從圖片中抽取所有結構化資訊，以 JSON 格式輸出。"}]},
        {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": "請從這張收據圖片中抽取所有資訊，以 JSON 格式輸出。"}]},
    ]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs = processor(text=[text], images=image_inputs, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=768, do_sample=False)
    generated_ids = output_ids[0][len(inputs["input_ids"][0]):]
    return processor.tokenizer.decode(generated_ids, skip_special_tokens=True)