dtflow 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

dtflow/converters.py

@@ -3,7 +3,8 @@
 
 提供与 HuggingFace datasets 等常用格式的互转功能。
 """
-from typing import List, Dict, Any, Optional, Union, Callable
+
+from typing import Any, Callable, Dict, List, Optional, Union
 
 
 def to_hf_dataset(data: List[Dict[str, Any]]):
@@ -44,7 +45,7 @@ def from_hf_dataset(dataset, split: Optional[str] = None) -> List[Dict[str, Any]
         >>> data = from_hf_dataset(my_dataset, split="train")
     """
     try:
-        from datasets import load_dataset, Dataset, DatasetDict
+        from datasets import Dataset, DatasetDict, load_dataset
     except ImportError:
         raise ImportError("需要安装 datasets: pip install datasets")
 
@@ -53,7 +54,7 @@ def from_hf_dataset(dataset, split: Optional[str] = None) -> List[Dict[str, Any]
         dataset = load_dataset(dataset, split=split)
 
     # 处理 DatasetDict
-    if hasattr(dataset, 'keys'):  # DatasetDict
+    if hasattr(dataset, "keys"):  # DatasetDict
         if split:
             dataset = dataset[split]
         else:
@@ -83,8 +84,9 @@ def to_hf_chat_format(
     Examples:
         >>> dt.transform(to_hf_chat_format())
     """
+
     def transform(item) -> dict:
-        messages = item.get(messages_field, []) if hasattr(item, 'get') else item[messages_field]
+        messages = item.get(messages_field, []) if hasattr(item, "get") else item[messages_field]
         result = {"messages": messages}
         if add_generation_prompt:
             result["add_generation_prompt"] = True
@@ -110,12 +112,14 @@ def from_openai_batch(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     for item in data:
         if item.get("response", {}).get("status_code") == 200:
             body = item["response"]["body"]
-            results.append({
-                "custom_id": item.get("custom_id"),
-                "content": body["choices"][0]["message"]["content"],
-                "model": body.get("model"),
-                "usage": body.get("usage"),
-            })
+            results.append(
+                {
+                    "custom_id": item.get("custom_id"),
+                    "content": body["choices"][0]["message"]["content"],
+                    "model": body.get("model"),
+                    "usage": body.get("usage"),
+                }
+            )
     return results
 
 
@@ -138,11 +142,12 @@ def to_openai_batch(
     Examples:
         >>> batch_input = dt.to(to_openai_batch(model="gpt-4o"))
     """
+
     def transform(item, idx=[0]) -> dict:
-        messages = item.get(messages_field, []) if hasattr(item, 'get') else item[messages_field]
+        messages = item.get(messages_field, []) if hasattr(item, "get") else item[messages_field]
 
         if custom_id_field:
-            custom_id = item.get(custom_id_field) if hasattr(item, 'get') else item[custom_id_field]
+            custom_id = item.get(custom_id_field) if hasattr(item, "get") else item[custom_id_field]
         else:
             custom_id = f"request-{idx[0]}"
             idx[0] += 1
@@ -154,7 +159,7 @@ def to_openai_batch(
             "body": {
                 "model": model,
                 "messages": messages,
-            }
+            },
         }
 
     return transform
@@ -189,8 +194,9 @@ def to_llama_factory(
     Returns:
         转换函数
     """
+
     def transform(item) -> dict:
-        get = lambda f: (item.get(f, "") if hasattr(item, 'get') else item.get(f, ""))
+        get = lambda f: (item.get(f, "") if hasattr(item, "get") else item.get(f, ""))
 
         result = {
             "instruction": get(instruction_field),
@@ -237,8 +243,13 @@ def to_axolotl(
     Returns:
         转换函数
     """
+
     def transform(item) -> dict:
-        conversations = item.get(conversations_field, []) if hasattr(item, 'get') else item.get(conversations_field, [])
+        conversations = (
+            item.get(conversations_field, [])
+            if hasattr(item, "get")
+            else item.get(conversations_field, [])
+        )
 
         # 如果已经是正确格式，直接返回
         if conversations and isinstance(conversations[0], dict):
@@ -246,11 +257,14 @@ def to_axolotl(
                 return {"conversations": conversations}
 
         # 尝试从 messages 格式转换
-        messages = item.get("messages", []) if hasattr(item, 'get') else item.get("messages", [])
+        messages = item.get("messages", []) if hasattr(item, "get") else item.get("messages", [])
         if messages:
             role_map = {"user": "human", "assistant": "gpt", "system": "system"}
             conversations = [
-                {from_key: role_map.get(m.get("role", ""), m.get("role", "")), value_key: m.get("content", "")}
+                {
+                    from_key: role_map.get(m.get("role", ""), m.get("role", "")),
+                    value_key: m.get("content", ""),
+                }
                 for m in messages
             ]
 
@@ -541,10 +555,12 @@ def to_swift_messages(
 
         for msg in messages:
             # 标准化格式
-            result_messages.append({
-                "role": msg.get("role", "user"),
-                "content": msg.get("content", ""),
-            })
+            result_messages.append(
+                {
+                    "role": msg.get("role", "user"),
+                    "content": msg.get("content", ""),
+                }
+            )
 
         return {"messages": result_messages}
 
@@ -749,8 +765,8 @@ def messages_to_text(
     fmt = templates[template]
 
     def transform(item) -> dict:
-        result = item.to_dict() if hasattr(item, 'to_dict') else dict(item)
-        messages = item.get(messages_field, []) if hasattr(item, 'get') else item[messages_field]
+        result = item.to_dict() if hasattr(item, "to_dict") else dict(item)
+        messages = item.get(messages_field, []) if hasattr(item, "get") else item[messages_field]
 
         parts = []
         for msg in messages:

dtflow/core.py

@@ -3,42 +3,32 @@ DataTransformer 核心模块
 
 专注于数据格式转换，提供简洁的 API。
 """
-from typing import List, Dict, Any, Optional, Callable, Union, Tuple, Literal
+
 from copy import deepcopy
 from dataclasses import dataclass
-import json
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
-from .storage.io import save_data, load_data
+import orjson
 
-# 尝试使用 orjson（更快的 JSON 序列化库）
-try:
-    import orjson
-    _HAS_ORJSON = True
-except ImportError:
-    _HAS_ORJSON = False
+from .lineage import LineageTracker
+from .storage.io import load_data, save_data
 
 
 def _fast_json_dumps(obj: Any) -> str:
-    """
-    快速 JSON 序列化，优先使用 orjson。
-
-    orjson 比标准 json 快约 10 倍，特别适合大量数据的序列化场景。
-    """
-    if _HAS_ORJSON:
-        # orjson.dumps 返回 bytes，需要 decode
-        return orjson.dumps(obj, option=orjson.OPT_SORT_KEYS).decode('utf-8')
-    else:
-        return json.dumps(obj, sort_keys=True, ensure_ascii=False)
+    """快速 JSON 序列化（使用 orjson，比标准 json 快约 10 倍）"""
+    return orjson.dumps(obj, option=orjson.OPT_SORT_KEYS).decode("utf-8")
 
 
 # ============ 错误处理 ============
 
+
 @dataclass
 class TransformError:
     """转换错误信息"""
-    index: int          # 原始数据索引
-    item: Dict          # 原始数据项
-    error: Exception    # 异常对象
+
+    index: int  # 原始数据索引
+    item: Dict  # 原始数据项
+    error: Exception  # 异常对象
 
     def __repr__(self) -> str:
         return f"TransformError(index={self.index}, error={self.error!r})"
@@ -61,9 +51,11 @@ class TransformErrors(Exception):
     def _build_message(self) -> str:
         if len(self.errors) == 1:
             return str(self.errors[0])
-        return f"转换失败 {len(self.errors)} 条记录:\n" + "\n".join(
-            f"  [{e.index}] {e.error}" for e in self.errors[:5]
-        ) + (f"\n  ... 还有 {len(self.errors) - 5} 条错误" if len(self.errors) > 5 else "")
+        return (
+            f"转换失败 {len(self.errors)} 条记录:\n"
+            + "\n".join(f"  [{e.index}] {e.error}" for e in self.errors[:5])
+            + (f"\n  ... 还有 {len(self.errors) - 5} 条错误" if len(self.errors) > 5 else "")
+        )
 
     def __iter__(self):
         return iter(self.errors)
@@ -102,8 +94,15 @@ class DataTransformer:
     - fields/stats: 数据信息
     """
 
-    def __init__(self, data: Optional[List[Dict[str, Any]]] = None):
+    def __init__(
+        self,
+        data: Optional[List[Dict[str, Any]]] = None,
+        _source_path: Optional[str] = None,
+        _lineage_tracker: Optional[LineageTracker] = None,
+    ):
         self._data = data if data is not None else []
+        self._source_path = _source_path
+        self._lineage_tracker = _lineage_tracker
 
     @property
     def data(self) -> List[Dict[str, Any]]:
@@ -122,23 +121,39 @@ class DataTransformer:
     # ============ 加载/保存 ============
 
     @classmethod
-    def load(cls, filepath: str) -> 'DataTransformer':
+    def load(cls, filepath: str, track_lineage: bool = False) -> "DataTransformer":
         """
         从文件加载数据。
 
         支持格式: jsonl, json, csv, parquet（自动检测）
+
+        Args:
+            filepath: 文件路径
+            track_lineage: 是否追踪血缘（默认 False）
         """
         data = load_data(filepath)
-        return cls(data)
+        tracker = LineageTracker(filepath) if track_lineage else None
+        return cls(data, _source_path=filepath, _lineage_tracker=tracker)
 
-    def save(self, filepath: str) -> None:
+    def save(self, filepath: str, lineage: bool = False) -> None:
         """
         保存数据到文件。
 
         支持格式: jsonl, json, csv, parquet（根据扩展名）
+
+        Args:
+            filepath: 文件路径
+            lineage: 是否保存血缘元数据（默认 False）
         """
         save_data(self._data, filepath)
 
+        # 保存血缘记录
+        if lineage and self._lineage_tracker:
+            lineage_path = self._lineage_tracker.save(filepath, len(self._data))
+            import sys
+
+            print(f"📜 血缘记录已保存: {lineage_path}", file=sys.stderr)
+
     # ============ 核心转换 ============
 
     def to(
@@ -215,7 +230,7 @@ class DataTransformer:
         func: Callable[[Any], Any],
         on_error: Literal["skip", "raise", "null"] = "skip",
         raw: bool = False,
-    ) -> 'DataTransformer':
+    ) -> "DataTransformer":
         """
         转换数据并返回新的 DataTransformer（支持链式调用）。
 
@@ -230,7 +245,16 @@ class DataTransformer:
             >>> # 原始模式（大数据集推荐）
             >>> dt.transform(lambda x: {"q": x["q"]}, raw=True).save("output.jsonl")
         """
-        return DataTransformer(self.to(func, on_error=on_error, raw=raw))
+        input_count = len(self._data)
+        result = self.to(func, on_error=on_error, raw=raw)
+        output_count = len(result)
+
+        # 传递血缘追踪器并记录操作
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("transform", {"func": func}, input_count, output_count)
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     # ============ 数据筛选 ============
 
@@ -239,7 +263,7 @@ class DataTransformer:
         func: Callable[[Any], bool],
         on_error: Literal["skip", "raise", "keep"] = "skip",
         raw: bool = False,
-    ) -> 'DataTransformer':
+    ) -> "DataTransformer":
         """
         筛选数据。
 
@@ -281,9 +305,14 @@ class DataTransformer:
         if errors:
             _print_error_summary(errors, len(self._data))
 
-        return DataTransformer(filtered)
+        # 传递血缘追踪器并记录操作
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("filter", {"func": func}, len(self._data), len(filtered))
 
-    def sample(self, n: int, seed: Optional[int] = None) -> 'DataTransformer':
+        return DataTransformer(filtered, _lineage_tracker=tracker)
+
+    def sample(self, n: int, seed: Optional[int] = None) -> "DataTransformer":
         """
         随机采样 n 条数据。
 
@@ -292,24 +321,39 @@ class DataTransformer:
             seed: 随机种子
         """
         import random
+
         if seed is not None:
             random.seed(seed)
 
+        input_count = len(self._data)
         data = self._data[:] if n >= len(self._data) else random.sample(self._data, n)
-        return DataTransformer(data)
 
-    def head(self, n: int = 10) -> 'DataTransformer':
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("sample", {"n": n, "seed": seed}, input_count, len(data))
+
+        return DataTransformer(data, _lineage_tracker=tracker)
+
+    def head(self, n: int = 10) -> "DataTransformer":
         """取前 n 条"""
-        return DataTransformer(self._data[:n])
+        data = self._data[:n]
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("head", {"n": n}, len(self._data), len(data))
+        return DataTransformer(data, _lineage_tracker=tracker)
 
-    def tail(self, n: int = 10) -> 'DataTransformer':
+    def tail(self, n: int = 10) -> "DataTransformer":
         """取后 n 条"""
-        return DataTransformer(self._data[-n:])
+        data = self._data[-n:]
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("tail", {"n": n}, len(self._data), len(data))
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def dedupe(
         self,
         key: Union[None, str, List[str], Callable[[Any], Any]] = None,
-    ) -> 'DataTransformer':
+    ) -> "DataTransformer":
         """
         数据去重。
 
@@ -338,7 +382,11 @@ class DataTransformer:
                 seen.add(k)
                 result.append(item)
 
-        return DataTransformer(result)
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("dedupe", {"key": key}, len(self._data), len(result))
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     def _get_dedupe_key(
         self,
@@ -367,7 +415,7 @@ class DataTransformer:
         threshold: float = 0.8,
         num_perm: int = 128,
         ngram: int = 3,
-    ) -> 'DataTransformer':
+    ) -> "DataTransformer":
         """
         基于 MinHash + LSH 的相似度去重。
 
@@ -388,9 +436,7 @@ class DataTransformer:
         try:
             from datasketch import MinHash, MinHashLSH
         except ImportError:
-            raise ImportError(
-                "相似度去重需要 datasketch 库，请安装: pip install datasketch"
-            )
+            raise ImportError("相似度去重需要 datasketch 库，请安装: pip install datasketch")
 
         if not self._data:
             return DataTransformer([])
@@ -400,10 +446,11 @@ class DataTransformer:
         # threshold=0.99 需要 num_perm>=512，threshold>=0.999 会需要极大的值(4096+)
         if threshold >= 0.999:
             import warnings
+
             warnings.warn(
                 f"阈值 {threshold} 过高，已自动调整为 0.99。"
                 f"如需更高精度，建议使用 dedupe() 精确去重。",
-                UserWarning
+                UserWarning,
             )
             threshold = 0.99
 
@@ -442,7 +489,17 @@ class DataTransformer:
 
         # 按原顺序保留数据
         result = [self._data[i] for i in sorted(keep_indices)]
-        return DataTransformer(result)
+
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record(
+                "dedupe_similar",
+                {"key": key, "threshold": threshold, "num_perm": num_perm, "ngram": ngram},
+                len(self._data),
+                len(result),
+            )
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     def _get_text_for_similarity(
         self,
@@ -457,14 +514,14 @@ class DataTransformer:
         else:
             raise ValueError(f"不支持的 key 类型: {type(key)}")
 
-    def _create_minhash(self, text: str, num_perm: int, ngram: int) -> 'MinHash':
+    def _create_minhash(self, text: str, num_perm: int, ngram: int) -> "MinHash":
         """创建文本的 MinHash 签名"""
         from datasketch import MinHash
 
         m = MinHash(num_perm=num_perm)
         # 使用字符级 n-gram（对中英文都适用）
         for i in range(len(text) - ngram + 1):
-            m.update(text[i:i + ngram].encode('utf-8'))
+            m.update(text[i : i + ngram].encode("utf-8"))
         return m
 
     # ============ 数据信息 ============
@@ -485,7 +542,7 @@ class DataTransformer:
 
         return sorted(all_fields)
 
-    def _extract_fields(self, obj: Any, prefix: str = '') -> List[str]:
+    def _extract_fields(self, obj: Any, prefix: str = "") -> List[str]:
         """递归提取字段名"""
         fields = []
         if isinstance(obj, dict):
@@ -516,25 +573,21 @@ class DataTransformer:
             field_stats[key] = {
                 "count": len(values),
                 "missing": len(self._data) - len(values),
-                "type": type(values[0]).__name__ if values else "unknown"
+                "type": type(values[0]).__name__ if values else "unknown",
             }
 
-        return {
-            "total": len(self._data),
-            "fields": sorted(all_keys),
-            "field_stats": field_stats
-        }
+        return {"total": len(self._data), "fields": sorted(all_keys), "field_stats": field_stats}
 
     # ============ 工具方法 ============
 
-    def copy(self) -> 'DataTransformer':
+    def copy(self) -> "DataTransformer":
         """深拷贝"""
         return DataTransformer(deepcopy(self._data))
 
     # ============ 数据合并 ============
 
     @classmethod
-    def concat(cls, *sources: Union[str, 'DataTransformer']) -> 'DataTransformer':
+    def concat(cls, *sources: Union[str, "DataTransformer"]) -> "DataTransformer":
         """
         拼接多个数据源。
 
@@ -564,7 +617,7 @@ class DataTransformer:
 
         return cls(all_data)
 
-    def __add__(self, other: Union[str, 'DataTransformer']) -> 'DataTransformer':
+    def __add__(self, other: Union[str, "DataTransformer"]) -> "DataTransformer":
         """
         使用 + 运算符拼接数据。
 
@@ -574,14 +627,20 @@ class DataTransformer:
         """
         return DataTransformer.concat(self, other)
 
-    def shuffle(self, seed: Optional[int] = None) -> 'DataTransformer':
+    def shuffle(self, seed: Optional[int] = None) -> "DataTransformer":
         """打乱顺序（返回新实例）"""
         import random
+
         data = self._data[:]
         if seed is not None:
             random.seed(seed)
         random.shuffle(data)
-        return DataTransformer(data)
+
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("shuffle", {"seed": seed}, len(self._data), len(data))
+
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def split(self, ratio: float = 0.8, seed: Optional[int] = None) -> tuple:
         """
@@ -596,7 +655,16 @@ class DataTransformer:
         """
         data = self.shuffle(seed).data
         split_idx = int(len(data) * ratio)
-        return DataTransformer(data[:split_idx]), DataTransformer(data[split_idx:])
+
+        # 分割后血缘追踪器各自独立
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("split", {"ratio": ratio, "seed": seed}, len(self._data), len(data))
+
+        return (
+            DataTransformer(data[:split_idx], _lineage_tracker=tracker),
+            DataTransformer(data[split_idx:], _lineage_tracker=tracker),
+        )
 
     # ============ 并行处理 ============
 
@@ -641,7 +709,7 @@ class DataTransformer:
         func: Callable[[Dict], bool],
         workers: Optional[int] = None,
         chunksize: int = 1000,
-    ) -> 'DataTransformer':
+    ) -> "DataTransformer":
         """
         并行执行过滤函数（使用多进程）。
 
@@ -700,18 +768,18 @@ class DictWrapper:
     """
 
     def __init__(self, data: Dict[str, Any]):
-        object.__setattr__(self, '_data', data)
+        object.__setattr__(self, "_data", data)
         # 构建规范化名称到原始名称的映射
         alias_map = {}
         for key in data.keys():
             sanitized = _sanitize_key(key)
             if sanitized != key:
                 alias_map[sanitized] = key
-        object.__setattr__(self, '_alias_map', alias_map)
+        object.__setattr__(self, "_alias_map", alias_map)
 
     def __getattr__(self, name: str) -> Any:
-        data = object.__getattribute__(self, '_data')
-        alias_map = object.__getattribute__(self, '_alias_map')
+        data = object.__getattribute__(self, "_data")
+        alias_map = object.__getattribute__(self, "_alias_map")
 
         # 先尝试直接匹配
         if name in data:
@@ -730,23 +798,23 @@ class DictWrapper:
         raise AttributeError(f"字段不存在: {name}")
 
     def __getitem__(self, key: str) -> Any:
-        data = object.__getattribute__(self, '_data')
+        data = object.__getattribute__(self, "_data")
         value = data[key]
         if isinstance(value, dict):
             return DictWrapper(value)
         return value
 
     def __contains__(self, key: str) -> bool:
-        data = object.__getattribute__(self, '_data')
+        data = object.__getattribute__(self, "_data")
         return key in data
 
     def __repr__(self) -> str:
-        data = object.__getattribute__(self, '_data')
+        data = object.__getattribute__(self, "_data")
         return repr(data)
 
     def get(self, key: str, default: Any = None) -> Any:
         """安全获取字段值"""
-        data = object.__getattribute__(self, '_data')
+        data = object.__getattribute__(self, "_data")
         value = data.get(key, default)
         if isinstance(value, dict):
             return DictWrapper(value)
@@ -754,4 +822,4 @@ class DictWrapper:
 
     def to_dict(self) -> Dict[str, Any]:
         """返回原始字典"""
-        return object.__getattribute__(self, '_data')
+        return object.__getattribute__(self, "_data")