Browse Source

fix(pdf_parse_union_core_v2): suppress FutureWarning from transformers

- Added warnings module to import list
- Implemented a warning catcher to ignore FutureWarning from the transformers module
- This change prevents unnecessary warning messages during model inference
myhloli 7 months ago
parent
commit
3cb156f549
1 changed files with 7 additions and 3 deletions
  1. 7 3
      magic_pdf/pdf_parse_union_core_v2.py

+ 7 - 3
magic_pdf/pdf_parse_union_core_v2.py

@@ -4,6 +4,7 @@ import os
 import re
 import statistics
 import time
+import warnings
 from typing import List
 
 import cv2
@@ -369,9 +370,12 @@ def do_predict(boxes: List[List[int]], model) -> List[int]:
     from magic_pdf.model.sub_modules.reading_oreder.layoutreader.helpers import (
         boxes2inputs, parse_logits, prepare_inputs)
 
-    inputs = boxes2inputs(boxes)
-    inputs = prepare_inputs(inputs, model)
-    logits = model(**inputs).logits.cpu().squeeze(0)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
+
+        inputs = boxes2inputs(boxes)
+        inputs = prepare_inputs(inputs, model)
+        logits = model(**inputs).logits.cpu().squeeze(0)
     return parse_logits(logits, len(boxes))