import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER = "your-username/billstructai-qwen-lora"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL, quantization_config=bnb_config, device_map="auto"
)
model = PeftModel.from_pretrained(base_model, ADAPTER)
model.eval()
SYSTEM_PROMPT = """You are a precise invoice information extraction assistant.
Your task is to extract structured invoice data from noisy OCR text.
Return only valid JSON.
Do not explain.
Do not add markdown.
Use null for missing fields.
Do not hallucinate values that are not present in the OCR text."""
ocr_text = "PASTE OCR TEXT HERE"
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": f"Extract invoice information from the OCR text below and return only valid JSON.\n\nOCR_TEXT:\n{ocr_text}"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
output = model.generate(**inputs, max_new_tokens=768, do_sample=False, pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id)
print(tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))