old code

app.py

@@ -43,18 +43,8 @@ tax_deductions = '''Extract the following information in the given format:
                 }
 '''
 
-
-
-
-def process_document(image):
-    # Save the uploaded image to a temporary file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
-        image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-        image.save(tmp_file.name)  # Save the image to the temporary file
-        image_path = tmp_file.name  # Get the path of the saved file
-
-
-        messages = [
+def demo(image_path, prompt):
+    messages = [
         {
             "role": "user",
             "content": [
@@ -62,18 +52,7 @@ def process_document(image):
                     "type": "image",
                     "image": image_path,  # Use the file path here
                 },
-                {"type": "text", "text":  '''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':""}}},
-                }
-'''},
+                {"type": "text", "text":  prompt},
             ],
         }
     ]
@@ -106,31 +85,40 @@ def process_document(image):
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )
     print("8")
-    print("output_text ---->>>>", output_text)
-    print("output text type --->>>> ", type(output_text))
-    print("First element", output_text[0])
 
     # Handle output text to convert it into JSON
-    json = str()
     try:
-        print("Trying ......")
         almost_json = output_text[0].split('\n')[-1].split('\n')[0]
         json = literal_eval(almost_json)
-    except Exception as e:
-        print(f"Got error {e}")
-        # json = output_text[0]  # Return raw output if JSON parsing fails
-        json = output_text
+    except:
+        json = output_text[0]  # Return raw output if JSON parsing fails
+    return json
+
+def process_document(image):
+    # Save the uploaded image to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
+        image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+        image.save(tmp_file.name)  # Save the image to the temporary file
+        image_path = tmp_file.name  # Get the path of the saved file
+
+    # 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
+    }
 
     # Optionally, you can delete the temporary file after use
     os.remove(image_path)
 
-    return json
+    return json_op
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=process_document,
     inputs="image",  # Gradio will handle the image input
-    outputs="text",
+    outputs="json",
     title="PaySlip_Demo_Model",
     examples=[["Slip_1.jpg"], ["Slip_2.jpg"]],
     cache_examples=False