dtflow 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

dtflow/converters.py

@@ -259,6 +259,462 @@ def to_axolotl(
     return transform
 
 
+def to_llama_factory_sharegpt(
+    messages_field: str = "messages",
+    system_field: Optional[str] = None,
+    tools_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 LLaMA-Factory ShareGPT 格式（多轮对话）。
+
+    输出格式:
+    {
+        "conversations": [
+            {"from": "human", "value": "..."},
+            {"from": "gpt", "value": "..."}
+        ],
+        "system": "...",      # 可选
+        "tools": "..."        # 可选
+    }
+
+    Args:
+        messages_field: 输入的 messages 字段名
+        system_field: 系统提示字段（如果为 None，从 messages 中提取）
+        tools_field: 工具描述字段
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_llama_factory_sharegpt())
+        >>> dt.transform(to_llama_factory_sharegpt(system_field="system_prompt"))
+    """
+    role_map = {
+        "user": "human",
+        "assistant": "gpt",
+        "system": "system",
+        "tool": "observation",
+        "function_call": "function_call",
+    }
+
+    def transform(item) -> dict:
+        get = lambda f: (item.get(f, "") if hasattr(item, "get") else item.get(f, ""))
+        messages = get(messages_field) or []
+
+        conversations = []
+        system_prompt = None
+
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+
+            # 提取 system 消息
+            if role == "system":
+                system_prompt = content
+                continue
+
+            mapped_role = role_map.get(role, role)
+            conversations.append({"from": mapped_role, "value": content})
+
+        result = {"conversations": conversations}
+
+        # 系统提示：优先使用指定字段，否则用从 messages 提取的
+        if system_field:
+            system = get(system_field)
+            if system:
+                result["system"] = system
+        elif system_prompt:
+            result["system"] = system_prompt
+
+        # 工具描述
+        if tools_field:
+            tools = get(tools_field)
+            if tools:
+                result["tools"] = tools
+
+        return result
+
+    return transform
+
+
+def to_llama_factory_vlm(
+    messages_field: str = "messages",
+    images_field: str = "images",
+    videos_field: Optional[str] = None,
+    system_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 LLaMA-Factory VLM（视觉语言模型）格式。
+
+    输出格式 (Alpaca 风格):
+    {
+        "instruction": "...",
+        "input": "",
+        "output": "...",
+        "images": ["path1.jpg", "path2.jpg"],  # 图片路径列表
+        "videos": ["path.mp4"],                 # 可选，视频路径列表
+        "system": "..."                         # 可选
+    }
+
+    Args:
+        messages_field: 输入的 messages 字段名
+        images_field: 图片路径字段名
+        videos_field: 视频路径字段名
+        system_field: 系统提示字段
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_llama_factory_vlm())
+        >>> dt.transform(to_llama_factory_vlm(images_field="image_paths"))
+    """
+
+    def transform(item) -> dict:
+        get = lambda f: item.get(f) if hasattr(item, "get") else item.get(f)
+        messages = get(messages_field) or []
+
+        instruction = ""
+        output = ""
+        system_prompt = None
+
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+
+            if role == "system":
+                system_prompt = content
+            elif role == "user":
+                instruction = content
+            elif role == "assistant":
+                output = content
+
+        result = {
+            "instruction": instruction,
+            "input": "",
+            "output": output,
+        }
+
+        # 图片
+        images = get(images_field)
+        if images:
+            result["images"] = images if isinstance(images, list) else [images]
+
+        # 视频
+        if videos_field:
+            videos = get(videos_field)
+            if videos:
+                result["videos"] = videos if isinstance(videos, list) else [videos]
+
+        # 系统提示
+        if system_field:
+            system = get(system_field)
+            if system:
+                result["system"] = system
+        elif system_prompt:
+            result["system"] = system_prompt
+
+        return result
+
+    return transform
+
+
+def to_llama_factory_vlm_sharegpt(
+    messages_field: str = "messages",
+    images_field: str = "images",
+    videos_field: Optional[str] = None,
+    system_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 LLaMA-Factory VLM ShareGPT 格式（多轮多模态对话）。
+
+    输出格式:
+    {
+        "conversations": [
+            {"from": "human", "value": "<image>描述这张图片"},
+            {"from": "gpt", "value": "这是一张..."}
+        ],
+        "images": ["path1.jpg"],
+        "system": "..."
+    }
+
+    Args:
+        messages_field: 输入的 messages 字段名
+        images_field: 图片路径字段名
+        videos_field: 视频路径字段名
+        system_field: 系统提示字段
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_llama_factory_vlm_sharegpt())
+    """
+    role_map = {"user": "human", "assistant": "gpt", "system": "system"}
+
+    def transform(item) -> dict:
+        get = lambda f: item.get(f) if hasattr(item, "get") else item.get(f)
+        messages = get(messages_field) or []
+
+        conversations = []
+        system_prompt = None
+
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+
+            if role == "system":
+                system_prompt = content
+                continue
+
+            mapped_role = role_map.get(role, role)
+            conversations.append({"from": mapped_role, "value": content})
+
+        result = {"conversations": conversations}
+
+        # 图片
+        images = get(images_field)
+        if images:
+            result["images"] = images if isinstance(images, list) else [images]
+
+        # 视频
+        if videos_field:
+            videos = get(videos_field)
+            if videos:
+                result["videos"] = videos if isinstance(videos, list) else [videos]
+
+        # 系统提示
+        if system_field:
+            system = get(system_field)
+            if system:
+                result["system"] = system
+        elif system_prompt:
+            result["system"] = system_prompt
+
+        return result
+
+    return transform
+
+
+# ============== ms-swift 格式转换器 ==============
+
+
+def to_swift_messages(
+    messages_field: str = "messages",
+    system_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 ms-swift messages 格式（标准格式）。
+
+    输出格式:
+    {
+        "messages": [
+            {"role": "system", "content": "..."},
+            {"role": "user", "content": "..."},
+            {"role": "assistant", "content": "..."}
+        ]
+    }
+
+    Args:
+        messages_field: 输入的 messages 字段名
+        system_field: 系统提示字段（如果需要额外添加）
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_swift_messages())
+    """
+
+    def transform(item) -> dict:
+        get = lambda f: item.get(f) if hasattr(item, "get") else item.get(f)
+        messages = get(messages_field) or []
+
+        # 复制 messages，避免修改原数据
+        result_messages = []
+
+        # 如果指定了 system_field，添加系统消息
+        if system_field:
+            system = get(system_field)
+            if system:
+                result_messages.append({"role": "system", "content": system})
+
+        for msg in messages:
+            # 标准化格式
+            result_messages.append({
+                "role": msg.get("role", "user"),
+                "content": msg.get("content", ""),
+            })
+
+        return {"messages": result_messages}
+
+    return transform
+
+
+def to_swift_query_response(
+    query_field: str = "query",
+    response_field: str = "response",
+    system_field: Optional[str] = None,
+    history_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 ms-swift query-response 格式。
+
+    输出格式:
+    {
+        "query": "用户问题",
+        "response": "模型回答",
+        "system": "系统提示",      # 可选
+        "history": [["q1", "r1"]]  # 可选
+    }
+
+    Args:
+        query_field: 用户问题字段
+        response_field: 模型回答字段
+        system_field: 系统提示字段
+        history_field: 历史对话字段
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_swift_query_response())
+        >>> # 从 messages 格式转换
+        >>> dt.transform(to_swift_query_response(query_field="messages"))
+    """
+
+    def transform(item) -> dict:
+        get = lambda f: item.get(f) if hasattr(item, "get") else item.get(f)
+
+        query = get(query_field)
+        response = get(response_field)
+
+        # 如果 query_field 是 messages，提取最后一轮对话
+        if isinstance(query, list):
+            messages = query
+            system_prompt = None
+            history = []
+            current_query = ""
+            current_response = ""
+
+            for i, msg in enumerate(messages):
+                role = msg.get("role", "")
+                content = msg.get("content", "")
+
+                if role == "system":
+                    system_prompt = content
+                elif role == "user":
+                    if current_query and current_response:
+                        history.append([current_query, current_response])
+                    current_query = content
+                    current_response = ""
+                elif role == "assistant":
+                    current_response = content
+
+            result = {
+                "query": current_query,
+                "response": current_response,
+            }
+
+            if system_prompt:
+                result["system"] = system_prompt
+            if history:
+                result["history"] = history
+
+            return result
+
+        # 直接使用字段
+        result = {
+            "query": query or "",
+            "response": response or "",
+        }
+
+        if system_field:
+            system = get(system_field)
+            if system:
+                result["system"] = system
+
+        if history_field:
+            history = get(history_field)
+            if history:
+                result["history"] = history
+
+        return result
+
+    return transform
+
+
+def to_swift_vlm(
+    messages_field: str = "messages",
+    images_field: str = "images",
+    videos_field: Optional[str] = None,
+    system_field: Optional[str] = None,
+) -> Callable:
+    """
+    转换为 ms-swift VLM（视觉语言模型）格式。
+
+    输出格式:
+    {
+        "messages": [
+            {"role": "user", "content": "<image>描述图片"},
+            {"role": "assistant", "content": "这是..."}
+        ],
+        "images": ["/path/to/image.jpg"]
+    }
+
+    Args:
+        messages_field: 输入的 messages 字段名
+        images_field: 图片路径字段名
+        videos_field: 视频路径字段名
+        system_field: 系统提示字段
+
+    Returns:
+        转换函数
+
+    Examples:
+        >>> dt.transform(to_swift_vlm())
+        >>> dt.transform(to_swift_vlm(images_field="image_paths"))
+    """
+
+    def transform(item) -> dict:
+        get = lambda f: item.get(f) if hasattr(item, "get") else item.get(f)
+        messages = get(messages_field) or []
+
+        result_messages = []
+
+        # 添加系统提示
+        if system_field:
+            system = get(system_field)
+            if system:
+                result_messages.append({"role": "system", "content": system})
+
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+
+            if role == "system" and not system_field:
+                result_messages.append({"role": "system", "content": content})
+            elif role in ("user", "assistant"):
+                result_messages.append({"role": role, "content": content})
+
+        result = {"messages": result_messages}
+
+        # 图片
+        images = get(images_field)
+        if images:
+            result["images"] = images if isinstance(images, list) else [images]
+
+        # 视频
+        if videos_field:
+            videos = get(videos_field)
+            if videos:
+                result["videos"] = videos if isinstance(videos, list) else [videos]
+
+        return result
+
+    return transform
+
+
 def messages_to_text(
     messages_field: str = "messages",
     output_field: str = "text",

dtflow/core.py

@@ -6,29 +6,16 @@ DataTransformer 核心模块
 from typing import List, Dict, Any, Optional, Callable, Union, Tuple, Literal
 from copy import deepcopy
 from dataclasses import dataclass
-import json
 
-from .storage.io import save_data, load_data
+import orjson
 
-# 尝试使用 orjson（更快的 JSON 序列化库）
-try:
-    import orjson
-    _HAS_ORJSON = True
-except ImportError:
-    _HAS_ORJSON = False
+from .storage.io import save_data, load_data
+from .lineage import LineageTracker
 
 
 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')
 
 
 # ============ 错误处理 ============
@@ -102,8 +89,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 +116,38 @@ 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(
@@ -230,7 +239,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)
 
     # ============ 数据筛选 ============
 
@@ -281,7 +299,12 @@ 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))
+
+        return DataTransformer(filtered, _lineage_tracker=tracker)
 
     def sample(self, n: int, seed: Optional[int] = None) -> 'DataTransformer':
         """
@@ -295,16 +318,30 @@ class DataTransformer:
         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)
+
+        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':
         """取后 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,
@@ -338,7 +375,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,
@@ -442,7 +483,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,
@@ -581,7 +632,12 @@ class DataTransformer:
         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 +652,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),
+        )
 
     # ============ 并行处理 ============