Commit 0a246f0f authored by myhloli's avatar myhloli
Browse files

refactor(magic_pdf): simplify device selection in model initialization

- Replace complex device selection logic with a single line using torch.device
- Remove redundant checks and imports for better readability and maintainability
parent 9b00f988
...@@ -338,17 +338,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang ...@@ -338,17 +338,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang
def model_init(model_name: str): def model_init(model_name: str):
from transformers import LayoutLMv3ForTokenClassification from transformers import LayoutLMv3ForTokenClassification
device = get_device() device = torch.device(get_device())
if torch.cuda.is_available():
device = torch.device('cuda')
elif str(device).startswith("npu"):
import torch_npu
if torch_npu.npu.is_available():
device = torch.device('npu')
else:
device = torch.device('cpu')
else:
device = torch.device('cpu')
if model_name == 'layoutreader': if model_name == 'layoutreader':
# 检测modelscope的缓存目录是否存在 # 检测modelscope的缓存目录是否存在
......
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