"vscode:/vscode.git/clone" did not exist on "885e3c201a9bc39a70a6bd2744c5ef6843a6c156"
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,6 +370,9 @@ def do_predict(boxes: List[List[int]], model) -> List[int]: ...@@ -369,6 +370,9 @@ 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)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
inputs = boxes2inputs(boxes) inputs = boxes2inputs(boxes)
inputs = prepare_inputs(inputs, model) inputs = prepare_inputs(inputs, model)
logits = model(**inputs).logits.cpu().squeeze(0) logits = model(**inputs).logits.cpu().squeeze(0)
......
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