wehub-resource-sync
128 lines
5.0 KiB
Python
128 lines
5.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import os
|
|
import platform
|
|
import re
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Literal, Optional, Tuple
|
|
|
|
from .dataset_meta import DATASET_MAPPING, DatasetMeta
|
|
|
|
_dataset_meta_mapping = None
|
|
|
|
|
|
@dataclass
|
|
class DatasetSyntax:
|
|
dataset: str
|
|
subsets: List[str] = field(default_factory=list)
|
|
dataset_sample: Optional[int] = None
|
|
use_hf: Optional[bool] = None
|
|
|
|
def __post_init__(self):
|
|
if os.path.isfile(self.dataset):
|
|
self.dataset_type = 'path'
|
|
else: # dataset_id or dataset_dir
|
|
self.dataset_type = 'repo'
|
|
|
|
def get_raw(self):
|
|
subsets = '/'.join(self.subsets)
|
|
dataset_sample = '' if self.dataset_sample is None else f'#{self.dataset_sample}'
|
|
return f'{self.dataset}{subsets}{dataset_sample}'
|
|
|
|
@staticmethod
|
|
def _safe_split(s: str,
|
|
sep: str,
|
|
use_0: bool,
|
|
split_mode: Literal['left', 'right'] = 'left') -> Tuple[Optional[str], Optional[str]]:
|
|
"""
|
|
use_0: When the length of the part is 1, is it considered as part0 or part1.
|
|
split_mode: use split or rsplit
|
|
"""
|
|
if s is None or len(s) == 0:
|
|
return None, None
|
|
if split_mode == 'left':
|
|
part = s.split(sep, 1)
|
|
else:
|
|
part = s.rsplit(sep, 1)
|
|
if len(part) == 1:
|
|
if use_0:
|
|
part = part[0], None
|
|
else:
|
|
part = None, part[0]
|
|
else:
|
|
assert len(part) == 2
|
|
return part
|
|
|
|
@classmethod
|
|
def parse(cls, dataset: str) -> 'DatasetSyntax':
|
|
"""Parse the dataset from the command line"""
|
|
# hf/ms::dataset_id or dataset_path:subset1/subset2/subset3#dataset_sample
|
|
if os.path.exists(dataset):
|
|
use_hf = None
|
|
else:
|
|
use_hf, dataset = cls._safe_split(dataset, '::', False)
|
|
if isinstance(use_hf, str):
|
|
use_hf = use_hf.lower()
|
|
use_hf = {'hf': True, 'ms': False}.get(use_hf)
|
|
if os.path.exists(dataset):
|
|
other, dataset_sample = dataset, None
|
|
else:
|
|
other, dataset_sample = cls._safe_split(dataset, '#', True, 'right')
|
|
if os.path.exists(other):
|
|
dataset, subsets = other, None
|
|
else:
|
|
dataset, subsets = cls._safe_split(other, ':', True)
|
|
|
|
if subsets is not None:
|
|
subsets = [subset.strip() for subset in subsets.split('/')]
|
|
if dataset_sample is not None:
|
|
dataset_sample = int(dataset_sample)
|
|
return cls(dataset.strip(), subsets or [], dataset_sample, use_hf)
|
|
|
|
def get_dataset_meta(self, use_hf: bool):
|
|
dataset_meta_mapping = self._get_dataset_meta_mapping()
|
|
dataset_type = self.dataset_type
|
|
if dataset_type == 'path':
|
|
dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset))
|
|
else:
|
|
dataset_type = 'repo' if os.path.isdir(self.dataset) else {True: 'hf', False: 'ms'}[use_hf]
|
|
dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset))
|
|
return dataset_meta or self._get_matched_dataset_meta(dataset_meta_mapping) or DatasetMeta()
|
|
|
|
@staticmethod
|
|
def _get_dataset_meta_mapping() -> Dict[Tuple[str, str], DatasetMeta]:
|
|
global _dataset_meta_mapping
|
|
if _dataset_meta_mapping is not None:
|
|
return _dataset_meta_mapping
|
|
_dataset_meta_mapping = {}
|
|
for dataset_meta in DATASET_MAPPING.values():
|
|
if dataset_meta.dataset_path is not None:
|
|
dataset_type = 'repo' if os.path.isdir(dataset_meta.dataset_path) else 'path'
|
|
_dataset_meta_mapping[(dataset_type, dataset_meta.dataset_path)] = dataset_meta
|
|
if dataset_meta.ms_dataset_id is not None:
|
|
_dataset_meta_mapping[('ms', dataset_meta.ms_dataset_id)] = dataset_meta
|
|
if dataset_meta.hf_dataset_id is not None:
|
|
_dataset_meta_mapping[('hf', dataset_meta.hf_dataset_id)] = dataset_meta
|
|
return _dataset_meta_mapping
|
|
|
|
@staticmethod
|
|
def get_dataset_name(dataset_id: str) -> str:
|
|
# compat hf hub
|
|
dataset_id = dataset_id.rstrip('/')
|
|
match_ = re.search('/datasets--.+?--(.+?)/snapshots/', dataset_id)
|
|
if match_ is not None:
|
|
return match_.group(1)
|
|
|
|
dataset_name = dataset_id.rsplit('/', 1)[-1]
|
|
if platform.system().lower() == 'windows':
|
|
dataset_name = dataset_name.rsplit('\\', 1)[-1]
|
|
return dataset_name
|
|
|
|
def _get_matched_dataset_meta(self, dataset_meta_mapping):
|
|
suffix_dataset_meta_mapping = {}
|
|
for dataset_name, dataset_meta in dataset_meta_mapping.items():
|
|
dataset_name = self.get_dataset_name(dataset_name[1])
|
|
suffix_dataset_meta_mapping[dataset_name] = dataset_meta
|
|
dataset_name = self.get_dataset_name(self.dataset)
|
|
dataset_meta = suffix_dataset_meta_mapping.get(dataset_name)
|
|
return dataset_meta
|