qwen update, per image, two output

app.py

@@ -1,46 +1,104 @@
 import re
 import gradio as gr
 
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
 import torch
-from transformers import DonutProcessor, VisionEncoderDecoderModel
+from ast import literal_eval
 
-processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
-model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+# 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)
 
-def process_document(image):
 
-    print(image)
-    print(f"Type of Image {image}")
-    # prepare encoder inputs
-    pixel_values = processor(image, return_tensors="pt").pixel_values
-    
-    # prepare decoder inputs
-    task_prompt = "<s_cord-v2>"
-    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-    
-    # generate answer
-    outputs = model.generate(
-        pixel_values.to(device),
-        decoder_input_ids=decoder_input_ids.to(device),
-        max_length=model.decoder.config.max_position_embeddings,
-        early_stopping=True,
-        pad_token_id=processor.tokenizer.pad_token_id,
-        eos_token_id=processor.tokenizer.eos_token_id,
-        use_cache=True,
-        num_beams=1,
-        bad_words_ids=[[processor.tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
+
+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",
     )
-    
-    # postprocess
-    sequence = processor.batch_decode(outputs.sequences)[0]
-    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
-    
-    return processor.token2json(sequence)
+    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)
+
+    return one, two
 
 # article = "<p style='text-align: center'><a href='https://www.xelpmoc.in/' target='_blank'>Made by Xelpmoc</a></p>"
 
@@ -51,7 +109,7 @@ demo = gr.Interface(
     title="Template-Free OCR model",
     # article=article,
     enable_queue=True,
-    examples=[["example.png"], ["example_2.png"], ["example_3.png"]],
+    examples=[["example.png"], ["example_2.png"]],
     cache_examples=False)
 
 demo.launch()