File size: 3,652 Bytes
87b2711
7c99e50
1ca62bc
87b2711
f52050f
7c99e50
f52050f
ea69d06
6082f24
f52050f
 
 
7c99e50
6082f24
f52050f
2fecabe
6082f24
f52050f
 
 
 
 
 
 
 
ec4bd6c
 
f52050f
 
 
 
 
 
 
 
 
 
 
ec4bd6c
f52050f
 
 
 
87b2711
f52050f
 
 
 
 
 
87b2711
f52050f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7c99e50
f52050f
 
 
 
 
 
 
 
 
 
6082f24
 
f52050f
6082f24
f52050f
 
6082f24
f52050f
 
 
87b2711
 
 
 
 
 
4222109
 
 
6082f24
4222109
7c99e50
6082f24
7c99e50
 
6082f24
7c99e50
4222109
6082f24
 
 
7c99e50
6082f24
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import os
import gradio as gr
import numpy as np
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
from ast import literal_eval

# Load the model on the available device(s)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
)

# Load the processor
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")

# Define your prompts
other_benifits = '''Extract the following information in the given format:
        {'other_benefits_and_information': {
            '401k eru: {'This Period':'', 'Year-to-Date':''}},
            'quota summary':
                            {
                            'sick:': '',
                            'vacation:': '',
                            }
            'payment method': 'eg. Direct payment',
            'Amount': 'eg. 12.99'
                            }
'''

tax_deductions = '''Extract the following information in the given format:
                {
                'tax_deductions': {
                    'federal:': {
                        'withholding tax:': {'Amount':'', 'Year-To_Date':""},
                        'ee social security tax:': {'Amount':'', 'Year-To_Date':""},
                        'ee medicare tax:': {'Amount':'', 'Year-To_Date':""}},
                    'california:': {
                        'withholding tax:': {'Amount':'', 'Year-To_Date':""},
                        'ee disability tax:': {'Amount':'', 'Year-To_Date':""}}},
                }
'''

def demo(image_path, prompt):
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,  # Use the file path here
                },
                {"type": "text", "text":  prompt},
            ],
        }
    ]

    # Preparation for inference
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    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=1500)
    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
    )

    # Handle output text to convert it into JSON
    try:
        almost_json = output_text[0].split('\n')[-1].split('\n')[0]
        json = literal_eval(almost_json)
    except:
        json = output_text[0]  # Return raw output if JSON parsing fails
    return json

def process_document(image):
    # Save the uploaded image temporarily and get its path
    image_path = image.name  # Gradio provides an interface to access the file name

    # Process the image with your model
    one = demo(image_path, other_benifits)
    two = demo(image_path, tax_deductions)
    json_op = {
        "tax_deductions": one,
        "other_benifits": two
    }
    return json_op

# Create Gradio interface
demo = gr.Interface(
    fn=process_document,
    inputs="image",  # Gradio will handle the image input
    outputs="json",
    title="PaySlip_Demo_Model",
    examples=[["Slip_1.jpg"], ["Slip_2.jpg"]],
    cache_examples=False
)

demo.launch()