Spaces:
Paused
Paused
| import re | |
| import gradio as gr | |
| from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor | |
| from qwen_vl_utils import process_vision_info | |
| import torch | |
| from ast import literal_eval | |
| # default: Load the model on the available device(s) | |
| model = Qwen2VLForConditionalGeneration.from_pretrained( | |
| "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto" | |
| ) | |
| # default processer | |
| processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| 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': '', | |
| 'Amount': '' | |
| } | |
| ''' | |
| 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_name, prompt): | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "image", | |
| "image": image_name, | |
| }, | |
| {"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 | |
| ) | |
| try: | |
| # almost_json = output_text[0].replace('```\n', '').replace('\n```', '') | |
| almost_json = output_text[0].split('```\n')[-1].split('\n```')[0] | |
| json = literal_eval(almost_json) | |
| except: | |
| try: | |
| # almost_json = output_text[0].replace('```json\n', '').replace('\n```', '') | |
| almost_json = output_text[0].split('```json\n')[-1].split('\n```')[0] | |
| json = literal_eval(almost_json) | |
| except: | |
| json = output_text[0] | |
| return json | |
| def process_document(image): | |
| one = demo(image, other_benifits) | |
| two = demo(image, tax_deductions) | |
| json_op = { | |
| "tax_deductions": one, | |
| "other_benifits": two | |
| } | |
| return json_op | |
| # article = "<p style='text-align: center'><a href='https://www.xelpmoc.in/' target='_blank'>Made by Xelpmoc</a></p>" | |
| demo = gr.Interface( | |
| fn=process_document, | |
| inputs="image", | |
| outputs="json", | |
| title="PaySlip_Demo_Model", | |
| # article=article, | |
| enable_queue=True, | |
| examples=[["Slip_1.jpg"], ["Slip_2.jpg"]], | |
| cache_examples=False) | |
| demo.launch() |