Browse Source

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
myhloli 8 months ago
parent
commit
0a246f0f40
1 changed files with 1 additions and 11 deletions
  1. 1 11
      magic_pdf/pdf_parse_union_core_v2.py

+ 1 - 11
magic_pdf/pdf_parse_union_core_v2.py

@@ -338,17 +338,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang
 
 def model_init(model_name: str):
     from transformers import LayoutLMv3ForTokenClassification
-    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')
+    device = torch.device(get_device())
 
     if model_name == 'layoutreader':
         # 检测modelscope的缓存目录是否存在