from transformers import AutoProcessor, AutoModelForMultimodalLM
from qwen_vl_utils import process_vision_info
from pathlib import Path
import torch

# Load QWEN
model = AutoModelForMultimodalLM.from_pretrained(
    "starride-teklia/DAI-NER-flat",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="auto",
)
processor = AutoProcessor.from_pretrained("starride-teklia/DAI-NER-flat")

# Prompt
SYSTEM = Path("system.txt").read_text()		# Available in the model directory
IMAGE = "record.jpg"		                     # https://europe.iiif.teklia.com/iiif/2/geneanet%2FArdennes_BMS%2F379692%2F00100.jpg/1425,793,1252,284/full/0/default.jpg
HTR_TRANSCRIPTION = "Mort de Pierre Soyer gardien de la redoutte du prez d'an...gne et ligne de cette frontiere\n\nL'an mil sept cens et neuf le vingt neuvième jour du mois de jeanvier a esté decedée en cette paroisse Pierre Soyer gardien au pre d'A[...]gne de la paroisse d'Aumont agee de quarante deux ans ou environs lequel a esté inhumé le trantième dudit mois dans le cimetier de cette paroisse ou nous lavons conduit avec les ceremonies accoustumée en présence avec les témoins qui ont signé avec nous\n\nJames Gileux T Stenva prb Mathieu Lallement."

messages = [
    {
        "role": "system",
        "content": [
           {
               "type": "text",
               "text": SYSTEM
           }
        ]
    },
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "record.jpg"
            },
            {
                "type": "text",
                "text": f"Convert this record in JSON. To help you, here is the transcription produced by an OCR model: {HTR_TRANSCRIPTION}"
            },
        ],
    }
]

# Preparation for inference
text = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=[text],
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = inputs.to("cuda")

# Inference: Generation of the output
generated_ids = model.generate(**inputs, max_new_tokens=1024)
generated_ids_trimmed = [
    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)[0]
print(output_text)

Output

{"événements.décès.année": "mil sept cens et neuf", "événements.décès.jour": "vingt neuvième", "événements.décès.mois": "janvier", "événements.décès.lieu": "en cette paroisse", "événements.acte.jour": "trantième", "événements.acte.mois": "dudit mois", "individus.défunt.prénom": "Pierre", "individus.défunt.nom": "Soyer", "individus.défunt.profession": "gardien au pre d'A[...]gne", "individus.défunt.âge": "quarante deux ans ou environs", "individus.défunt.paroisse": "d'Aumont", "individus.officiant.nom": "Mathieu Lallement"}