Selaa lähdekoodia

修改模型下载和数据集下载方式

lxylxy123321 1 viikko sitten
vanhempi
sitoutus
45d3544ac4

+ 6 - 3
backend/app/services/dataset_service.py

@@ -107,11 +107,14 @@ def _download_modelscope_dataset(dataset_id: str) -> tuple[Path, Path, int]:
     """用 MsDataset 下载并转为 JSONL。"""
     from modelscope.msdatasets import MsDataset
 
+    # 先尝试加载完整数据集
     try:
-        ds = MsDataset.load(dataset_id, subset_name="default", split="train")
-    except Exception:
-        # 回退:不带参数,自动选择第一个 split
         ds = MsDataset.load(dataset_id)
+    except Exception:
+        try:
+            ds = MsDataset.load(dataset_id, subset_name="default", split="train")
+        except Exception:
+            ds = MsDataset.load(dataset_id, split="train")
 
     ds_dir = settings.processed_dir / f"ms_{dataset_id.replace('/', '_')}"
     ds_dir.mkdir(parents=True, exist_ok=True)

+ 6 - 2
backend/app/services/model_service.py

@@ -42,9 +42,13 @@ async def download_model(model_id: str, use_modelscope: bool = False) -> dict[st
     """从 HF 或 ModelScope 下载模型到本地缓存。"""
     try:
         if use_modelscope:
-            from modelscope import snapshot_download as ms_download
+            from modelscope.hub.snapshot_download import snapshot_download as ms_snapshot_download
 
-            local_path = ms_download(model_id, cache_dir=str(settings.models_dir))
+            download_dir = str(settings.models_dir / model_id.replace("/", "_"))
+            local_path = ms_snapshot_download(
+                model_id,
+                local_dir=download_dir,
+            )
         else:
             from huggingface_hub import snapshot_download