avatar63/qwen-receipt-extractor

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import json

BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTER = "avatar63/qwen-receipt-extractor"

INSTRUCTION = (
    "Extract the following fields from the OCR text as JSON: "
    "company_name, address, date, total_amount, line_items "
    "(each with item_name, quantity, price). "
    "Use null for any field that cannot be determined."
)

tokenizer = AutoTokenizer.from_pretrained(ADAPTER, trust_remote_code=True)
base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, dtype=torch.float16, device_map="auto", trust_remote_code=True
)
model = PeftModel.from_pretrained(base, ADAPTER)
model.eval()

noisy_text = """
RELI4NCE FR3SH
Sh0p N0 12, 5ect0r 18
D4te: O5-ll-2O24
Net P4y4ble: 34O.OO
"""

messages = [
    {"role": "system", "content": INSTRUCTION},
    {"role": "user", "content": noisy_text}
]

text = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.1,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )

generated = outputs[0][inputs["input_ids"].shape[1]:]
result = tokenizer.decode(generated, skip_special_tokens=True)
print(json.loads(result))