dtflow 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

dtflow/parallel.py

@@ -0,0 +1,115 @@
+"""
+并行处理模块
+
+提供多进程并行处理工具，用于加速大数据集的 token 统计和 schema 验证。
+"""
+
+from multiprocessing import Pool, cpu_count
+from typing import Callable, List, Optional, TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def parallel_map(
+    func: Callable[[T], R],
+    data: List[T],
+    workers: Optional[int] = None,
+    threshold: int = 1000,
+    chunksize: Optional[int] = None,
+) -> List[R]:
+    """
+    并行 map 操作。
+
+    Args:
+        func: 处理函数（必须可 pickle，不能是 lambda 或闭包）
+        data: 数据列表
+        workers: 进程数，None 则使用 CPU 核数
+        threshold: 数据量阈值，低于此值使用串行
+        chunksize: 每个进程的任务块大小，None 则自动计算
+
+    Returns:
+        处理结果列表（保持顺序）
+    """
+    n = len(data)
+
+    # 数据量小或指定单进程，使用串行
+    if n < threshold or workers == 1:
+        return [func(item) for item in data]
+
+    workers = workers or cpu_count()
+    workers = min(workers, n)  # 进程数不超过数据量
+
+    # 自动计算 chunksize
+    if chunksize is None:
+        chunksize = max(1, n // (workers * 4))
+
+    with Pool(processes=workers) as pool:
+        return pool.map(func, data, chunksize=chunksize)
+
+
+def parallel_imap(
+    func: Callable[[T], R],
+    data: List[T],
+    workers: Optional[int] = None,
+    threshold: int = 1000,
+    chunksize: Optional[int] = None,
+):
+    """
+    并行 imap 操作（惰性迭代器版本，支持进度回调）。
+
+    Args:
+        func: 处理函数（必须可 pickle）
+        data: 数据列表
+        workers: 进程数，None 则使用 CPU 核数
+        threshold: 数据量阈值，低于此值使用串行
+        chunksize: 每个进程的任务块大小
+
+    Yields:
+        处理结果（按顺序）
+    """
+    n = len(data)
+
+    # 数据量小或指定单进程，使用串行
+    if n < threshold or workers == 1:
+        for item in data:
+            yield func(item)
+        return
+
+    workers = workers or cpu_count()
+    workers = min(workers, n)
+
+    if chunksize is None:
+        chunksize = max(1, n // (workers * 4))
+
+    with Pool(processes=workers) as pool:
+        for result in pool.imap(func, data, chunksize=chunksize):
+            yield result
+
+
+def get_optimal_workers(data_size: int, default: Optional[int] = None) -> int:
+    """
+    根据数据量计算最优进程数。
+
+    Args:
+        data_size: 数据量
+        default: 用户指定的进程数，None 则自动计算
+
+    Returns:
+        最优进程数
+    """
+    if default is not None:
+        return default
+
+    cpu_cores = cpu_count()
+
+    # 数据量小于阈值，单进程
+    if data_size < 1000:
+        return 1
+
+    # 数据量适中，使用一半 CPU
+    if data_size < 10000:
+        return max(1, cpu_cores // 2)
+
+    # 大数据量，使用全部 CPU
+    return cpu_cores

dtflow/schema.py

@@ -26,10 +26,35 @@ Schema 验证模块
     results = dt.validate_schema(schema)
 """
 
-from dataclasses import dataclass, field as dataclass_field
-from typing import Any, Callable, Dict, List, Literal, Optional, Set, Union
+from dataclasses import dataclass
+from dataclasses import field as dataclass_field
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+
+from .utils.field_path import _parse_path, get_field
+
+
+def _validate_item_wrapper(args: tuple) -> Tuple[int, bool, list]:
+    """
+    验证单条数据（用于多进程）。
+
+    Args:
+        args: (index, item, schema_fields) 元组
+
+    Returns:
+        (index, is_valid, errors_as_dicts) - 返回字典列表而非对象（pickle 兼容）
+    """
+    idx, item, fields = args
+    # 在子进程中重建 Schema
+    schema = Schema(fields)
+    result = schema.validate(item)
+
+    if result.valid:
+        return (idx, True, [])
+    else:
+        # 将错误转换为字典（pickle 兼容）
+        errors = [{"path": e.path, "message": e.message, "value": e.value} for e in result.errors]
+        return (idx, False, errors)
 
-from .utils.field_path import get_field, _parse_path, _get_value_by_segments
 
 # 支持的类型
 FieldType = Literal["str", "int", "float", "bool", "list", "dict", "any"]
@@ -162,9 +187,7 @@ class Field:
 
         # 选项检查
         if self.choices is not None and value not in self.choices:
-            errors.append(
-                ValidationError(path, f"值必须是 {self.choices} 之一", value)
-            )
+            errors.append(ValidationError(path, f"值必须是 {self.choices} 之一", value))
 
         # 正则表达式检查
         if self.pattern is not None and isinstance(value, str):
@@ -324,9 +347,7 @@ class Schema:
 
         return errors
 
-    def validate_batch(
-        self, data: List[dict], max_errors: int = 100
-    ) -> List[tuple]:
+    def validate_batch(self, data: List[dict], max_errors: int = 100) -> List[tuple]:
         """
         批量验证数据
 
@@ -350,9 +371,76 @@ class Schema:
 
         return failed
 
+    def validate_parallel(
+        self,
+        data: List[dict],
+        workers: Optional[int] = None,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> tuple:
+        """
+        并行验证数据列表。
+
+        Args:
+            data: 数据列表
+            workers: 进程数，None 自动检测，1 禁用并行
+            progress_callback: 进度回调函数
+
+        Returns:
+            (valid_data, invalid_indices_results) 元组
+            - valid_data: 有效数据列表
+            - invalid_indices_results: [(index, ValidationResult), ...] 无效数据
+        """
+        if not data:
+            return [], []
+
+        total = len(data)
+        use_parallel = workers != 1 and total >= 1000
+
+        valid_data = []
+        invalid_results = []
+
+        if use_parallel:
+            from .parallel import get_optimal_workers, parallel_imap
+
+            actual_workers = get_optimal_workers(total, workers)
+            # 准备参数：(index, item, schema_fields)
+            args_list = [(i, item, self._fields) for i, item in enumerate(data)]
+
+            for i, (idx, is_valid, result_data) in enumerate(
+                parallel_imap(
+                    _validate_item_wrapper,
+                    args_list,
+                    workers=actual_workers,
+                    threshold=1000,
+                )
+            ):
+                if is_valid:
+                    valid_data.append(data[idx])
+                else:
+                    # 重建 ValidationResult（因为不能直接 pickle）
+                    errors = [
+                        ValidationError(path=e["path"], message=e["message"], value=e.get("value"))
+                        for e in result_data
+                    ]
+                    invalid_results.append((idx, ValidationResult(valid=False, errors=errors)))
+                if progress_callback:
+                    progress_callback(i + 1, total)
+        else:
+            # 串行处理
+            for i, item in enumerate(data):
+                result = self.validate(item)
+                if result.valid:
+                    valid_data.append(item)
+                else:
+                    invalid_results.append((i, result))
+                if progress_callback:
+                    progress_callback(i + 1, total)
+
+        return valid_data, invalid_results
+
     def __repr__(self) -> str:
         field_strs = [f"  {path}: {field_def}" for path, field_def in self._fields.items()]
-        return f"Schema({{\n" + ",\n".join(field_strs) + "\n}})"
+        return "Schema({\n" + ",\n".join(field_strs) + "\n}})"
 
 
 # ============================================================================
@@ -461,9 +549,7 @@ def sharegpt_schema(
     """
     return Schema(
         {
-            "conversations": Field(
-                type="list", required=True, min_length=min_conversations
-            ),
+            "conversations": Field(type="list", required=True, min_length=min_conversations),
             "conversations[*].from": Field(
                 type="str", required=True, choices=[human_role, gpt_role]
             ),

dtflow/tokenizers.py

@@ -122,8 +122,8 @@ def _get_tiktoken_encoder(model: str):
                 _tokenizer_cache[model] = tiktoken.get_encoding(model)
             else:
                 _tokenizer_cache[model] = tiktoken.encoding_for_model(model)
-        except ImportError:
-            raise ImportError("需要安装 tiktoken: pip install tiktoken")
+        except ImportError as e:
+            raise ImportError("需要安装 tiktoken: pip install tiktoken") from e
     return _tokenizer_cache[model]
 
 
@@ -149,12 +149,12 @@ def _get_hf_tokenizer(model: str):
 
                 tokenizer = AutoTokenizer.from_pretrained(resolved, trust_remote_code=True)
                 _tokenizer_cache[resolved] = ("transformers", tokenizer)
-            except ImportError:
+            except ImportError as e:
                 raise ImportError(
                     "需要安装 tokenizers 或 transformers:\n"
                     "  pip install tokenizers huggingface_hub  (推荐，更轻量)\n"
                     "  pip install transformers"
-                )
+                ) from e
     return _tokenizer_cache[resolved]
 
 
@@ -309,12 +309,29 @@ def _std(counts: List[int], avg: float) -> float:
     return variance**0.5
 
 
+def _count_item_tokens(args: tuple) -> int:
+    """
+    计算单条数据的 token 数（用于多进程）。
+
+    Args:
+        args: (item, fields, model, backend) 元组
+    """
+    item, fields, model, backend = args
+    total = 0
+    for field in fields:
+        value = get_field_with_spec(item, field, default="")
+        if value:
+            total += count_tokens(str(value), model=model, backend=backend)
+    return total
+
+
 def token_stats(
     data: List[Dict[str, Any]],
     fields: Union[str, List[str]],
     model: str = DEFAULT_MODEL,
     backend: Optional[str] = None,
     progress_callback: Optional[Callable[[int, int], None]] = None,
+    workers: Optional[int] = None,
 ) -> Dict[str, Any]:
     """
     统计数据集的 token 信息。
@@ -325,6 +342,7 @@ def token_stats(
         model: 模型名称或别名，如 "qwen2.5", "gpt-4" 等
         backend: 后端选择，None 则自动检测
         progress_callback: 进度回调函数，接收 (current, total) 两个参数
+        workers: 进程数，None 自动检测，1 表示禁用并行
 
     Returns:
         统计信息字典，包含:
@@ -342,17 +360,42 @@ def token_stats(
     if not data:
         return {"total_tokens": 0, "count": 0}
 
-    counts = []
     total_items = len(data)
-    for i, item in enumerate(data):
-        total = 0
-        for field in fields:
-            value = get_field_with_spec(item, field, default="")
-            if value:
-                total += count_tokens(str(value), model=model, backend=backend)
-        counts.append(total)
-        if progress_callback:
-            progress_callback(i + 1, total_items)
+    _backend = backend or _auto_backend(model)
+
+    # 判断是否使用多进程
+    use_parallel = workers != 1 and total_items >= 1000
+
+    if use_parallel:
+        from .parallel import get_optimal_workers, parallel_imap
+
+        actual_workers = get_optimal_workers(total_items, workers)
+        # 准备参数
+        args_list = [(item, fields, model, _backend) for item in data]
+        counts = []
+        for i, result in enumerate(
+            parallel_imap(
+                _count_item_tokens,
+                args_list,
+                workers=actual_workers,
+                threshold=1000,
+            )
+        ):
+            counts.append(result)
+            if progress_callback:
+                progress_callback(i + 1, total_items)
+    else:
+        # 串行处理
+        counts = []
+        for i, item in enumerate(data):
+            total = 0
+            for field in fields:
+                value = get_field_with_spec(item, field, default="")
+                if value:
+                    total += count_tokens(str(value), model=model, backend=_backend)
+            counts.append(total)
+            if progress_callback:
+                progress_callback(i + 1, total_items)
 
     sorted_counts = sorted(counts)
     avg = sum(counts) / len(counts)
@@ -548,12 +591,27 @@ def messages_token_filter(
     return filter_func
 
 
+def _count_messages_tokens_wrapper(args: tuple) -> Optional[Dict[str, int]]:
+    """
+    计算单条 messages 的 token 数（用于多进程）。
+
+    Args:
+        args: (item, messages_field, model, backend) 元组
+    """
+    item, messages_field, model, backend = args
+    messages = get_field_with_spec(item, messages_field, default=[])
+    if messages:
+        return _count_messages_tokens(messages, model=model, backend=backend)
+    return None
+
+
 def messages_token_stats(
     data: List[Dict[str, Any]],
     messages_field: str = "messages",
     model: str = DEFAULT_MODEL,
     backend: Optional[str] = None,
     progress_callback: Optional[Callable[[int, int], None]] = None,
+    workers: Optional[int] = None,
 ) -> Dict[str, Any]:
     """
     统计数据集中 messages 的 token 信息。
@@ -564,6 +622,7 @@ def messages_token_stats(
         model: 模型名称或别名
         backend: 后端，None 则自动检测
         progress_callback: 进度回调函数，接收 (current, total) 两个参数
+        workers: 进程数，None 自动检测，1 表示禁用并行
 
     Returns:
         统计信息字典，包含:
@@ -581,14 +640,38 @@ def messages_token_stats(
     if not data:
         return {"count": 0, "total_tokens": 0}
 
-    all_stats = []
     total_items = len(data)
-    for i, item in enumerate(data):
-        messages = get_field_with_spec(item, messages_field, default=[])
-        if messages:
-            all_stats.append(_count_messages_tokens(messages, model=model, backend=_backend))
-        if progress_callback:
-            progress_callback(i + 1, total_items)
+
+    # 判断是否使用多进程
+    use_parallel = workers != 1 and total_items >= 1000
+
+    all_stats = []
+    if use_parallel:
+        from .parallel import get_optimal_workers, parallel_imap
+
+        actual_workers = get_optimal_workers(total_items, workers)
+        args_list = [(item, messages_field, model, _backend) for item in data]
+
+        for i, result in enumerate(
+            parallel_imap(
+                _count_messages_tokens_wrapper,
+                args_list,
+                workers=actual_workers,
+                threshold=1000,
+            )
+        ):
+            if result is not None:
+                all_stats.append(result)
+            if progress_callback:
+                progress_callback(i + 1, total_items)
+    else:
+        # 串行处理
+        for i, item in enumerate(data):
+            messages = get_field_with_spec(item, messages_field, default=[])
+            if messages:
+                all_stats.append(_count_messages_tokens(messages, model=model, backend=_backend))
+            if progress_callback:
+                progress_callback(i + 1, total_items)
 
     if not all_stats:
         return {"count": 0, "total_tokens": 0}

{dtflow-0.5.7.dist-info → dtflow-0.5.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dtflow
-Version: 0.5.7
+Version: 0.5.8
 Summary: A flexible data transformation tool for ML training formats (SFT, RLHF, Pretrain)
 Project-URL: Homepage, https://github.com/yourusername/DataTransformer
 Project-URL: Documentation, https://github.com/yourusername/DataTransformer#readme
@@ -452,6 +452,7 @@ dt run pipeline.yaml --input=new_data.jsonl --output=result.jsonl
 dt token-stats data.jsonl --field=messages --model=gpt-4
 dt token-stats data.jsonl --field=messages[-1].content   # 统计最后一条消息
 dt token-stats data.jsonl --field=text --detailed
+dt token-stats data.jsonl --workers=4                    # 多进程加速（数据量大时自动启用）
 
 # 数据对比
 dt diff v1/train.jsonl v2/train.jsonl
@@ -480,7 +481,11 @@ dt dedupe data.jsonl --key=text --similar=0.8   # 相似度去重
 dt concat a.jsonl b.jsonl -o merged.jsonl
 
 # 数据统计
-dt stats data.jsonl
+dt stats data.jsonl                                       # 快速模式
+dt stats data.jsonl --full                                # 完整模式（含值分布）
+dt stats data.jsonl --full --field=category               # 指定字段统计
+dt stats data.jsonl --full --expand=tags                  # 展开 list 字段统计元素分布
+dt stats data.jsonl --full --expand='messages[*].role'    # 展开嵌套 list 字段
 
 # Claude Code Skill 安装
 dt install-skill                              # 安装到 ~/.claude/skills/
@@ -491,6 +496,7 @@ dt validate data.jsonl --preset=openai_chat           # 使用预设 schema 验
 dt validate data.jsonl --preset=alpaca --verbose      # 详细输出
 dt validate data.jsonl --preset=sharegpt --filter-invalid -o valid.jsonl  # 过滤出有效数据
 dt validate data.jsonl --preset=dpo --max-errors=100  # 限制错误输出数量
+dt validate data.jsonl --preset=openai_chat --workers=4  # 多进程加速
 ```
 
 ### 字段路径语法

{dtflow-0.5.7.dist-info → dtflow-0.5.8.dist-info}/RECORD

@@ -1,17 +1,18 @@
-dtflow/SKILL.md,sha256=sHf6i6DKUCca5zvSJ67VHu05tFlST4mYgnoURXVe1g0,7836
-dtflow/__init__.py,sha256=Ee7CDDxzki69MEGeXB5bczuMts5OwZZ-jVsKjH_rD_0,3031
-dtflow/__main__.py,sha256=3LXTku09Fw1dsgTUtX1UJCmE20qKeZpNga3UqmI3UiY,12145
+dtflow/SKILL.md,sha256=nh12TTq_eRzl5O2CTgsiS809BBVR49kmpZ8n7UprMHI,9552
+dtflow/__init__.py,sha256=tofhUr_PMnsONnB3Hu-mwUrD4Q3bV7Kw_0S6dQw6ig8,3031
+dtflow/__main__.py,sha256=p8oZKQhwq04shCB3y_pkXjf-SZ4PZvg5PXdyUP-5rYA,13497
 dtflow/converters.py,sha256=X3qeFD7FCOMnfiP3MicL5MXimOm4XUYBs5pczIkudU0,22331
 dtflow/core.py,sha256=qMo6B3LK--TWRK7ZBKObGcs3pKFnd0NPoaM0T8JC7Jw,38135
 dtflow/framework.py,sha256=jyICi_RWHjX7WfsXdSbWmP1SL7y1OWSPyd5G5Y-lvg4,17578
 dtflow/lineage.py,sha256=jie3OL1qK90-_cOOqqLbhSJ1oGUktDM1x5HRpQ5Qiyc,12800
+dtflow/parallel.py,sha256=EnIdGEGMrZUNT2-CBIV93UFfpqr_jU_heqqvdGXcP-Y,3046
 dtflow/pipeline.py,sha256=zZaC4fg5vsp_30Fhbg75vu0yggsdvf28bWBiVDWzZ6Y,13901
 dtflow/presets.py,sha256=qa8WQJhbNMuGxqqgA9BFadEBwDB9s0zWNxxhzF3q1K8,4701
-dtflow/schema.py,sha256=IFcij22_UFKcgKT1YWwRg2QJO0vcAvCb1arZmsGByts,16824
+dtflow/schema.py,sha256=zCZNEAqTMT1BS_p2t0CYczR5S9rqyDREa7ZsYI5pFGA,19885
 dtflow/streaming.py,sha256=dxpNd1-Wz_PTLTdvM5qn06_2TJr5NRlIIuw0LOSS2Iw,24755
-dtflow/tokenizers.py,sha256=7ZAelSmcDxLWH5kICgH9Q1ULH3_BfDZb9suHMjJJRZU,20589
+dtflow/tokenizers.py,sha256=GFQsuLSLn2GHn2kaXhJkP8G85lgsdLzYtJNbppQhYPE,23408
 dtflow/cli/__init__.py,sha256=QhZ-thgx9IBTFII7T_hdoWFUl0CCsdGQHN5ZEZw2XB0,423
-dtflow/cli/clean.py,sha256=y9VCRibgK1j8WIY3h0XZX0m93EdELQC7TdnseMWwS-0,17799
+dtflow/cli/clean.py,sha256=KuE9ODjD9gSZUIHaD2mQLTDO-1PDwN7EqUpj8EQfVCs,25663
 dtflow/cli/commands.py,sha256=zKUG-B9Az-spqyqM00cR8Sgc2UgeOPQDThJFHWDNO_w,1336
 dtflow/cli/common.py,sha256=gCwnF5Sw2ploqfZJO_z3Ms9mR1HNT7Lj6ydHn0uVaIw,13817
 dtflow/cli/io_ops.py,sha256=BMDisP6dxzzmSjYwmeFwaHmpHHPqirmXAWeNTD-9MQM,13254
@@ -19,16 +20,16 @@ dtflow/cli/lineage.py,sha256=_lNh35nF9AA0Zy6FyZ4g8IzrXH2ZQnp3inF-o2Hs1pw,1383
 dtflow/cli/pipeline.py,sha256=QNEo-BJlaC1CVnVeRZr7TwfuZYloJ4TebIzJ5ALzry0,1426
 dtflow/cli/sample.py,sha256=pubpx4AIzsarBEalD150MC2apYQSt4bal70IZkTfFO0,15475
 dtflow/cli/skill.py,sha256=opiTEBejA7JHKrEMftMOPDQlOgZ4n59rwaHXGU1Nukk,2022
-dtflow/cli/stats.py,sha256=u4ehCfgw1X8WuOyAjrApMRgcIO3BVmINbsTjxEscQro,24086
+dtflow/cli/stats.py,sha256=HkTZD80h4tzYXTtMnfpjLUMP6kl_es6ifcmExxzGdMU,31813
 dtflow/cli/transform.py,sha256=w6xqMOxPxQvL2u_BPCfpDHuPSC9gmcqMPVN8s-B6bbY,15052
-dtflow/cli/validate.py,sha256=65aGVlMS_Rq0Ch0YQ-TclVJ03RQP4CnG137wthzb8Ao,4384
+dtflow/cli/validate.py,sha256=Frs-jKcDHmYozpmIYZueDSX5o2i1Xn-WW81FGUyUrng,5796
 dtflow/storage/__init__.py,sha256=C0jpWNQU808Ezz7lWneddABal3wILy8ijFUNiSKbHV4,362
 dtflow/storage/io.py,sha256=ZH2aSE-S89gpy3z4oTqhcqWf4u10OdkDoyul7o_YBDI,23374
 dtflow/utils/__init__.py,sha256=Pn-ltwV04fBQmeZG7FxInDQmzH29LYOi90LgeLMEuQk,506
 dtflow/utils/display.py,sha256=OeOdTh6mbDwSkDWlmkjfpTjy2QG8ZUaYU0NpHUWkpEQ,5881
 dtflow/utils/field_path.py,sha256=K8nU196RxTSJ1OoieTWGcYOWl9KjGq2iSxCAkfjECuM,7621
 dtflow/utils/helpers.py,sha256=JXN176_B2pm53GLVyZ1wj3wrmBJG52Tkw6AMQSdj7M8,791
-dtflow-0.5.7.dist-info/METADATA,sha256=mlWaRHSM1ZucQrAa8PGcHzjHj2RQPBynnmdA_JoNSNI,23899
-dtflow-0.5.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dtflow-0.5.7.dist-info/entry_points.txt,sha256=dadIDOK7Iu9pMxnMPBfpb4aAPe4hQbBOshpQYjVYpGc,44
-dtflow-0.5.7.dist-info/RECORD,,
+dtflow-0.5.8.dist-info/METADATA,sha256=Tm_dfdQfGlShyDt95fNQ87JXiBRnf6mfDgx827h3Rnc,24487
+dtflow-0.5.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dtflow-0.5.8.dist-info/entry_points.txt,sha256=dadIDOK7Iu9pMxnMPBfpb4aAPe4hQbBOshpQYjVYpGc,44
+dtflow-0.5.8.dist-info/RECORD,,