Views
No views yet

🚀 Note: This model is intended for learning purposes only and not for production applications. The extracted structured data may vary based on project needs.
Qwen/Qwen2.5-VL-3B-Instruct1%pip install -q "transformers>=4.49.0" accelerate datasets "qwen-vl-utils[decord]==0.0.8"
2
3import os
4import PIL
5import torch
6from qwen_vl_utils import process_vision_info
7from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
8
9
10MODEL_ID="zackriya/diagram2graph-adapters"
11MAX_PIXELS = 1280 * 28 * 28
12MIN_PIXELS = 256 * 28 * 28
13
14
15model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
16 MODEL_ID,
17 device_map="auto",
18 torch_dtype=torch.bfloat16
19)
20
21processor = Qwen2_5_VLProcessor.from_pretrained(
22 MODEL_ID,
23 min_pixels=MIN_PIXELS,
24 max_pixels=MAX_PIXELS
25)
26
27
28SYSTEM_MESSAGE = """You are a Vision Language Model specialized in extracting structured data from visual representations of process and flow diagrams.
29Your task is to analyze the provided image of a diagram and extract the relevant information into a well-structured JSON format.
30The diagram includes details such as nodes and edges. each of them have their own attributes.
31Focus on identifying key data fields and ensuring the output adheres to the requested JSON structure.
32Provide only the JSON output based on the extracted information. Avoid additional explanations or comments."""
33
34def run_inference(image):
35 messages= [
36 {
37 "role": "system",
38 "content": [{"type": "text", "text": SYSTEM_MESSAGE}],
39 },
40 {
41 "role": "user",
42 "content": [
43 {
44 "type": "image",
45 # this image is handled by qwen_vl_utils's process_visio_Info so no need to worry about pil image or path
46 "image": image,
47 },
48 {
49 "type": "text",
50 "text": "Extract data in JSON format, Only give the JSON",
51 },
52 ],
53 },
54 ]
55
56 text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
57 image_inputs, _ = process_vision_info(messages)
58
59 inputs = processor(
60 text=[text],
61 images=image_inputs,
62 return_tensors="pt",
63 )
64 inputs = inputs.to('cuda')
65
66 generated_ids = model.generate(**inputs, max_new_tokens=512)
67 generated_ids_trimmed = [
68 out_ids[len(in_ids):]
69 for in_ids, out_ids
70 in zip(inputs.input_ids, generated_ids)
71 ]
72
73 output_text = processor.batch_decode(
74 generated_ids_trimmed,
75 skip_special_tokens=True,
76 clean_up_tokenization_spaces=False
77 )
78 return output_text
79image = eval_dataset[9]['image'] # PIL image
80# `image` could be URL or relative path to the image
81output = run_inference(image)
82
83# JSON loading
84import json
85json.loads(output[0])| Samples | (Base)Node F1 | (Fine)Node F1 | (Base)Edge F1 | (Fine)Edge F1 |
|---|---|---|---|---|
| image_sample_1 | 0.46 | 1.0 | 0.59 | 0.71 |
| image_sample_2 | 0.67 | 0.57 | 0.25 | 0.25 |
| image_sample_3 | 1.0 | 1.0 | 0.25 | 0.75 |
| image_sample_4 | 0.5 | 0.83 | 0.15 | 0.62 |
| image_sample_5 | 0.72 | 0.78 | 0.0 | 0.48 |
| image_sample_6 | 0.75 | 0.75 | 0.29 | 0.67 |
| image_sample_7 | 0.6 | 1.0 | 1.0 | 1.0 |
| image_sample_8 | 0.6 | 1.0 | 1.0 | 1.0 |
| image_sample_9 | 1.0 | 1.0 | 0.55 | 0.77 |
| image_sample_10 | 0.67 | 0.8 | 0.0 | 1.0 |
| image_sample_11 | 0.8 | 0.8 | 0.5 | 1.0 |
| image_sample_12 | 0.67 | 1.0 | 0.62 | 0.75 |
| image_sample_13 | 1.0 | 1.0 | 0.73 | 0.67 |
| image_sample_14 | 0.74 | 0.95 | 0.56 | 0.67 |
| image_sample_15 | 0.86 | 0.71 | 0.67 | 0.67 |
| image_sample_16 | 0.75 | 1.0 | 0.8 | 0.75 |
| image_sample_17 | 0.8 | 1.0 | 0.63 | 0.73 |
| image_sample_18 | 0.83 | 0.83 | 0.33 | 0.43 |
| image_sample_19 | 0.75 | 0.8 | 0.06 | 0.22 |
| image_sample_20 | 0.81 | 1.0 | 0.23 | 0.75 |
| Mean | 0.749 | 0.891 | 0.4605 | 0.6945 |