Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2025-12-16 16:27:45 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1253 from modelscope/fix/import

Browse Source 
fix import
This commit is contained in:
Xingjun.Wang
2025-03-11 10:53:12 +08:00
committed by GitHub
parent 890a792569 c45bc6d19f
commit f0f9bdc331
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
modelscope/utils/hf_util/pipeline_builder.py
View File

@@ -1,13 +1,8 @@
import os
from typing import Optional, Union
import torch
from transformers import Pipeline as PipelineHF
from transformers import PreTrainedModel, TFPreTrainedModel, pipeline
from transformers.pipelines import check_task, get_task
from modelscope.hub import snapshot_download
from modelscope.utils.hf_util.patcher import _patch_pretrained_class, patch_hub
from modelscope.utils.hf_util.patcher import _patch_pretrained_class
def _get_hf_device(device):
@@ -21,6 +16,7 @@ def _get_hf_device(device):
def _get_hf_pipeline_class(task, model):
from transformers.pipelines import check_task, get_task
if not task:
task = get_task(model)
normalized_task, targeted_task, task_options = check_task(task)
@@ -35,7 +31,9 @@ def hf_pipeline(
framework: Optional[str] = None,
device: Optional[Union[int, str, 'torch.device']] = None,
**kwargs,
) -> PipelineHF:
) -> 'transformers.Pipeline':
from transformers import pipeline
if isinstance(model, str):
if not os.path.exists(model):
model = snapshot_download(model)