Commit 3cb156f5 authored by myhloli's avatar myhloli
Browse files

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
parent 59d6b195
...@@ -4,6 +4,7 @@ import os ...@@ -4,6 +4,7 @@ import os
import re import re
import statistics import statistics
import time import time
import warnings
from typing import List from typing import List
import cv2 import cv2
...@@ -369,9 +370,12 @@ def do_predict(boxes: List[List[int]], model) -> List[int]: ...@@ -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 ( from magic_pdf.model.sub_modules.reading_oreder.layoutreader.helpers import (
boxes2inputs, parse_logits, prepare_inputs) boxes2inputs, parse_logits, prepare_inputs)
inputs = boxes2inputs(boxes) with warnings.catch_warnings():
inputs = prepare_inputs(inputs, model) warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
logits = model(**inputs).logits.cpu().squeeze(0)
inputs = boxes2inputs(boxes)
inputs = prepare_inputs(inputs, model)
logits = model(**inputs).logits.cpu().squeeze(0)
return parse_logits(logits, len(boxes)) return parse_logits(logits, len(boxes))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment