Update ai_mapping.py

ai_mapping.py

@@ -39,14 +39,38 @@ def extract_key_values_with_layoutlm(text_data: str, pdf_path: str) -> Dict[str,
             pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))  # 300 DPI
             img_path = f"{pdf_path}_page_{page_num}.png"
             pix.save(img_path)
-            image = Image.open(img_path)
+            image = Image.open(img_path).convert("RGB")
 
-            encoding = feature_extractor(images=[image], text=text_data.splitlines(), return_tensors="pt")
+            # Tokenize text
+            words = text_data.splitlines()
+            encoding = tokenizer(
+                words,
+                return_tensors="pt",
+                truncation=True,
+                padding=True,
+                max_length=512
+            )
             input_ids = encoding["input_ids"]
             attention_mask = encoding["attention_mask"]
 
+            # Process image to get bounding boxes
+            image_encoding = feature_extractor(image, return_tensors="pt")
+            bbox = image_encoding["bbox"][0]  # Shape: (num_tokens, 4)
+
+            # Ensure bbox length matches input_ids
+            if len(bbox) < len(input_ids[0]):
+                bbox = torch.cat([bbox, torch.zeros((len(input_ids[0]) - len(bbox), 4), dtype=torch.int64)])
+            elif len(bbox) > len(input_ids[0]):
+                bbox = bbox[:len(input_ids[0])]
+
+            # Pass inputs to the model
             with torch.no_grad():
-                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+                outputs = model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    bbox=bbox,
+                    pixel_values=image_encoding["pixel_values"]
+                )
                 predictions = torch.argmax(outputs.logits, dim=2)
 
             tokens = tokenizer.convert_ids_to_tokens(input_ids[0])