distflow 1.0.0__py3-none-any.whl

distflow/cache/protocol.py

@@ -0,0 +1,7 @@
+from typing import Any, Protocol
+
+
+class CacheProtocol(Protocol):
+    async def load_cache(self, cache_key: str) -> dict[str, Any] | None: ...
+
+    async def save_cache(self, cache_key: str, cache_value: dict[str, Any]) -> bool: ...

distflow/cache/redis_cache.py

@@ -0,0 +1,122 @@
+import asyncio
+import json
+from typing import Any
+
+from redis.asyncio import Redis
+
+from distflow.utils import logger
+
+
+class RedisCache:
+    """使用 Redis 作为缓存后端的实现.
+
+    通过 Redis 客户端直接与 Redis 服务通信，实现分布式缓存。
+    使用 semaphore 限制并发请求数量。
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://127.0.0.1:6379",
+        max_concurrent_requests: int = 50,
+        redis_db: int = 0,
+    ) -> None:
+        """初始化Redis缓存.
+
+        Args:
+            redis_url: Redis 连接 URL，例如 "redis://127.0.0.1:6379"
+            max_concurrent_requests: 最大并发请求数
+            redis_db: Redis 数据库编号，默认为 0
+        """
+        self._semaphore = asyncio.Semaphore(max_concurrent_requests)
+
+        # 初始化 Redis 客户端
+        self._redis: Redis | None = None
+        self._redis_url = redis_url
+        self._redis_db = redis_db
+
+    def _get_redis(self) -> Redis:
+        """获取或创建 Redis 客户端."""
+        if self._redis is None:
+            self._redis = Redis.from_url(
+                self._redis_url,
+                db=self._redis_db,
+                decode_responses=True,
+            )
+            try:
+                # 测试连接
+                self._redis.ping()
+                logger.info(
+                    f"成功连接到 Redis: {self._redis_url}, DB: {self._redis_db}"
+                )
+            except Exception as e:
+                logger.error(
+                    f"无法连接到 Redis: {self._redis_url}, DB: {self._redis_db}, 错误: {e}"
+                )
+                raise ConnectionError(
+                    f"无法连接到 Redis: {self._redis_url}, DB: {self._redis_db}"
+                ) from e
+        return self._redis
+
+    async def load_cache(self, cache_key: str) -> dict[str, Any] | None:
+        """从 Redis 获取单个缓存值（受 semaphore 限制并发）.
+
+        Args:
+            cache_key: 缓存键
+
+        Returns:
+            缓存值字典，如果不存在则返回 None
+        """
+        for attempt in range(3):
+            async with self._semaphore:
+                try:
+                    redis = self._get_redis()
+                    cached_data = await redis.get(cache_key)
+                    if cached_data:
+                        return json.loads(cached_data)
+                    return None
+                except Exception as e:
+                    logger.warning(
+                        f"Redis 缓存查询失败 {attempt + 1} / 3: {type(e).__name__}: {e}"
+                    )
+                    await asyncio.sleep(0.1 * (attempt + 1))  # 简单的指数退避
+                    self._redis = None  # 重置 Redis 客户端以尝试重新连接
+        return None
+
+    async def save_cache(self, cache_key: str, cache_value: dict[str, Any]) -> bool:
+        """设置单个缓存值到 Redis（受 semaphore 限制并发）.
+
+        Args:
+            cache_key: 缓存键
+            cache_value: 缓存值
+
+        Returns:
+            是否成功
+        """
+        for attempt in range(3):
+            async with self._semaphore:
+                try:
+                    redis = self._get_redis()
+                    serialized = json.dumps(cache_value)
+                    await redis.set(cache_key, serialized)
+                    return True
+                except Exception as e:
+                    logger.warning(
+                        f"Redis 缓存写入失败 {attempt + 1} / 3: {type(e).__name__}: {e}"
+                    )
+                    await asyncio.sleep(0.1 * (attempt + 1))  # 简单的指数退避
+                    self._redis = None  # 重置 Redis 客户端以尝试重新连接
+        return False
+
+    async def close(self) -> None:
+        """关闭 Redis 连接."""
+        if self._redis:
+            await self._redis.close()
+            logger.info("Redis 连接已关闭")
+
+    async def __aenter__(self) -> "RedisCache":
+        """异步上下文管理器入口."""
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """异步上下文管理器退出."""
+        await self.close()

distflow/data/data_formatter.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import Any, Protocol, cast, runtime_checkable
+
+from distflow.data.types import DatasetProcessOutputItem, MessageData
+
+
+@runtime_checkable
+class FormatterProtocol(Protocol):
+    def format(self, raw_item: dict[str, Any]) -> DatasetProcessOutputItem: ...
+
+
+class AlpacaFormatter:
+    def __init__(self, *, user_key: str, assistant_key: str) -> None:
+        self.user_key = user_key
+        self.assistant_key = assistant_key
+
+    def format(self, raw_item: dict[str, Any]) -> DatasetProcessOutputItem:
+        assert (
+            self.user_key in raw_item
+        ), f"User key '{self.user_key}' not found in raw item"
+        assert (
+            self.assistant_key in raw_item
+        ), f"Assistant key '{self.assistant_key}' not found in raw item"
+        user_content = raw_item[self.user_key]
+        assert isinstance(
+            user_content, str
+        ), f"User content must be a string, got {type(user_content).__name__}: {user_content}"
+        assistant_content = raw_item[self.assistant_key]
+        assert isinstance(
+            assistant_content, str
+        ), f"Assistant content must be a string, got {type(assistant_content).__name__}: {assistant_content}"
+
+        return DatasetProcessOutputItem(
+            messages=[
+                cast(MessageData, {"role": "user", "content": user_content}),
+                cast(MessageData, {"role": "assistant", "content": assistant_content}),
+            ],
+            meta={
+                "user_key": self.user_key,
+                "assistant_key": self.assistant_key,
+                "raw_item": raw_item,
+            },
+        )
+
+
+class ShareGptFormatter:
+    def __init__(self, *, conversations_key: str) -> None:
+        self.conversations_key = conversations_key
+
+    def format(self, raw_item: dict[str, Any]) -> DatasetProcessOutputItem:
+        assert (
+            self.conversations_key in raw_item
+        ), f"Conversations key '{self.conversations_key}' not found in raw item"
+        conversations = raw_item[self.conversations_key]
+        assert isinstance(
+            conversations, list
+        ), f"Conversations must be a list, got {type(conversations).__name__}: {conversations}"
+
+        messages: list[MessageData] = []
+        for conv in conversations:
+            if isinstance(conv, dict):
+                role = conv.get("role")
+                content = conv.get("content")
+                if role is not None and content is not None:
+                    messages.append(
+                        cast(MessageData, {"role": role, "content": content})
+                    )
+
+        return DatasetProcessOutputItem(
+            messages=messages,
+            meta={"conversations_key": self.conversations_key, "raw_item": raw_item},
+        )

distflow/data/data_loader.py

@@ -0,0 +1,64 @@
+import builtins
+import random
+from typing import Any, Literal, cast
+
+from distflow.data.data_formatter import FormatterProtocol
+from distflow.data.types import DatasetProcessOutputItem
+from distflow.utils import logger
+
+
+def load_dataset(
+    dataset_name: str,
+    data_path: str,
+    load_type: Literal["datasets", "modelscope", "pandas"],
+    formatter: FormatterProtocol,
+    data_size: int = -1,
+    split: str = "train",
+    sep: str = "\t",
+    dtype: str = "str",
+    shuffle_seed: int = 42,
+    use_json: bool = False,
+) -> tuple[str, list[DatasetProcessOutputItem]]:
+    logger.info(f"开始加载数据集: {dataset_name}, 路径: {data_path}, 类型: {load_type}")
+
+    # 数据大小
+    logger.debug(f"数据大小限制: {data_size if data_size > 0 else '全部'}")
+
+    match load_type:
+        case "datasets":
+            from datasets import load_dataset
+
+            logger.debug(f"使用 datasets 加载, split={split}, use_json={use_json}")
+            if use_json:
+                dataset = load_dataset("json", data_files=data_path, split=split)
+            else:
+                dataset = load_dataset(path=data_path, split=split)
+        case "modelscope":
+            from modelscope.msdatasets import MsDataset
+
+            logger.debug(f"使用 modelscope 加载, split={split}")
+            dataset = MsDataset.load(data_path, split=split)
+        case "pandas":
+            from datasets import Dataset, load_dataset
+            from pandas import read_csv
+
+            logger.debug("使用 pandas 加载")
+            dtype_actual = getattr(builtins, dtype)
+            df = read_csv(data_path, sep=sep, dtype=dtype_actual)
+            dataset = Dataset.from_pandas(df)
+
+    logger.info(f"数据集加载完成，总样本数: {len(dataset)}")
+
+    random.seed(shuffle_seed)
+    logger.debug(f"使用随机种子: {shuffle_seed}")
+    random_indices = list(range(len(dataset)))
+    if data_size > 0 and data_size < len(dataset):
+        logger.info(f"随机采样 {data_size} 条数据")
+        random_indices = random.sample(random_indices, data_size)
+    else:
+        logger.info("使用全部数据")
+        random.shuffle(random_indices)
+    sampled_data = cast(list[dict[str, Any]], [dataset[i] for i in random_indices])
+    logger.debug(f"采样完成，开始格式化数据")
+    formatted_data = [formatter.format(data_item) for data_item in sampled_data]
+    return dataset_name, formatted_data

distflow/data/types.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class MessageData(BaseModel):  # type: ignore[misc]
+    role: str
+    content: str | dict[str, Any]
+
+
+class DatasetProcessOutputItem(BaseModel):  # type: ignore[misc]
+    messages: list[MessageData]
+    meta: dict[str, Any]

distflow/embed/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from distflow.embed.types import EmbeddingInputItem, EmbeddingResult
+
+
+class BaseEmbed(ABC):
+    def __init__(self, model_name: str) -> None:
+        self.model_name = model_name
+
+    @abstractmethod
+    def embed(self, dataset: list[EmbeddingInputItem]) -> list[EmbeddingResult]:
+        """异步嵌入计算.
+
+        Args:
+            dataset: 待嵌入的数据项列表
+
+        Returns:
+            嵌入结果列表
+        """

distflow/embed/cache_wrapper.py

@@ -0,0 +1,154 @@
+import asyncio
+import hashlib
+import json
+from collections.abc import Coroutine
+from typing import Any
+
+from distflow.cache.protocol import CacheProtocol
+from distflow.embed.base import BaseEmbed
+from distflow.embed.types import EmbeddingInputItem, EmbeddingResult
+from distflow.utils import logger
+
+
+def dict_to_hash(d: dict[Any, Any]) -> str:
+    """生成字典的SHA256哈希摘要"""
+    s = json.dumps(d, sort_keys=True).encode()
+    return hashlib.sha256(s).hexdigest()
+
+
+class CachedEmbed(BaseEmbed):
+    """使用 Redis 作为缓存后端的嵌入包装器.
+
+    通过 RedisCache 类与 Redis 服务通信，实现分布式缓存。
+    """
+
+    def __init__(
+        self,
+        embedder: BaseEmbed,
+        cache: CacheProtocol,
+        cache_model_id: str | None = None,
+        legacy_key: bool = False,
+    ) -> None:
+        """初始化缓存嵌入器.
+
+        Args:
+            embedder: 底层嵌入器，用于计算未缓存的数据
+            cache: 符合 CacheProtocol 的缓存实现
+            cache_model_id: 用于缓存键的模型标识符，默认为模型路径。
+                            可用于在移动模型后仍使用旧缓存。
+            legacy_key: 是否使用旧版缓存键格式（包含完整 data_item），
+                       默认为 False（使用新版：仅 model_id + messages）
+        """
+        self.embedder = embedder
+        self._cache = cache
+        self.model_path = (
+            getattr(embedder, "model_name", None)
+            or getattr(embedder, "model_path", None)
+            or "unknown"
+        )
+        # 用于缓存键的模型标识符
+        self.cache_model_id = cache_model_id if cache_model_id else self.model_path
+        self.legacy_key = legacy_key
+
+        super().__init__(self.model_path)
+
+    def _build_cache_key(self, item: EmbeddingInputItem) -> str:
+        """构建缓存键.
+
+        Args:
+            item: 输入数据项
+
+        Returns:
+            SHA256 哈希键
+        """
+        if self.legacy_key:
+            # 旧版格式：包含完整 data_item（包含 messages 和 meta）
+            key_payload = {
+                "model_path": self.model_path,
+                "data_item": item.model_dump(),
+            }
+        else:
+            # 新版格式：仅使用 cache_model_id 和 messages（不含 meta）
+            key_payload = {
+                "model_id": self.cache_model_id,
+                "messages": [msg.model_dump() for msg in item.messages],
+            }
+        return dict_to_hash(key_payload)
+
+    def embed(self, dataset: list[EmbeddingInputItem]) -> list[EmbeddingResult]:
+        """异步执行嵌入计算，使用 Redis 缓存.
+
+        Args:
+            dataset: 待嵌入的数据项列表
+
+        Returns:
+            嵌入结果列表
+        """
+        logger.info(f"开始缓存嵌入计算，数据量: {len(dataset)}")
+
+        # 并发查询所有缓存
+        cache_keys = [self._build_cache_key(item) for item in dataset]
+        cache_tasks = [self._cache.load_cache(key) for key in cache_keys]
+
+        async def _run_all_get_cache() -> list[dict[str, Any] | None | BaseException]:
+            return await asyncio.gather(*cache_tasks, return_exceptions=True)
+
+        cached_values = asyncio.run(_run_all_get_cache())
+
+        # 分离缓存命中和未命中的项
+        results: list[EmbeddingResult | None] = [None] * len(dataset)
+        missing_items: list[EmbeddingInputItem] = []
+        missing_indices: list[int] = []
+        missing_keys: list[str] = []
+
+        for idx, (item, key, cached_result) in enumerate(
+            zip(dataset, cache_keys, cached_values)
+        ):
+            # 处理异常结果
+            if isinstance(cached_result, BaseException):
+                logger.debug(f"缓存查询异常，将重新计算: {cached_result}")
+                missing_items.append(item)
+                missing_indices.append(idx)
+                missing_keys.append(key)
+            elif cached_result is None:
+                missing_items.append(item)
+                missing_indices.append(idx)
+                missing_keys.append(key)
+            else:
+                results[idx] = EmbeddingResult(
+                    embedding=cached_result["embedding"],
+                    data_item=item,
+                    meta=cached_result.get("meta", item.meta),
+                )
+                logger.debug(f"缓存命中: {key[:16]}...")
+
+        logger.info(f"缓存命中: {len(dataset) - len(missing_items)}/{len(dataset)}")
+
+        # 计算未缓存的嵌入
+        if missing_items:
+            new_results = self.embedder.embed(missing_items)
+
+            # 并发写入缓存
+            write_tasks: list[Coroutine[Any, Any, bool]] = []
+            for key, idx, result in zip(missing_keys, missing_indices, new_results):
+                cache_value = {
+                    "embedding": result.embedding,
+                    "meta": result.meta,
+                }
+                write_tasks.append(self._cache.save_cache(key, cache_value))
+                results[idx] = EmbeddingResult(
+                    embedding=result.embedding,
+                    data_item=dataset[idx],
+                    meta=result.meta,
+                )
+
+            # 等待所有写入完成
+            async def _run_all_save_cache() -> list[bool | BaseException]:
+                return await asyncio.gather(*write_tasks, return_exceptions=False)
+
+            write_results = asyncio.run(_run_all_save_cache())
+            success_count = sum(1 for r in write_results if r is True)
+            logger.info(f"缓存写入完成: {success_count}/{len(write_tasks)} 成功")
+
+        logger.info(f"嵌入计算完成，共 {len(results)} 条结果")
+        return [result for result in results if result is not None]

distflow/embed/sentence_transformers.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, override
+
+from distflow.data.types import MessageData
+from distflow.embed.base import BaseEmbed
+from distflow.embed.types import EmbeddingInputItem, EmbeddingResult
+from distflow.utils import logger
+from distflow.utils.timing import timing_context
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+
+class SentenceTransformersEmbed(BaseEmbed):
+    """基于 sentence-transformers 的嵌入器实现.
+
+    使用 sentence-transformers 库进行文本嵌入计算，支持批量处理和归一化。
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        device: str = "cuda",
+        batch_size: int = 32,
+        normalize_embeddings: bool = True,
+        trust_remote_code: bool = False,
+        prompt: str | None = None,
+    ) -> None:
+        """初始化 SentenceTransformersEmbed.
+
+        Args:
+            model_name: 模型名称或路径
+            device: 运行设备，默认为 "cuda"
+            batch_size: 批处理大小，默认为 32
+            normalize_embeddings: 是否对嵌入向量进行归一化，默认为 True
+            trust_remote_code: 是否信任远程代码，默认为 False
+            prompt: 可选的前缀提示文本，会添加到每个输入文本前面
+        """
+        logger.info(f"创建 SentenceTransformersEmbed，模型: {model_name}")
+        # 存储配置用于延迟初始化
+        self._model_name = model_name
+        self._device = device
+        self._batch_size = batch_size
+        self._normalize_embeddings = normalize_embeddings
+        self._trust_remote_code = trust_remote_code
+        self._prompt = prompt
+        self._model: SentenceTransformer | None = None
+        super().__init__(model_name)
+
+    def _ensure_initialized(self) -> None:
+        """延迟初始化模型 - 仅在需要嵌入计算时才调用."""
+        if self._model is not None:
+            return
+
+        logger.info(f"开始加载 Sentence Transformers 模型: {self._model_name}")
+        logger.debug(
+            f"配置参数: device={self._device}, batch_size={self._batch_size}, "
+            f"normalize_embeddings={self._normalize_embeddings}"
+        )
+
+        with timing_context("模型加载"):
+            from sentence_transformers import SentenceTransformer
+
+            self._model = SentenceTransformer(
+                self._model_name,
+                device=self._device,
+                trust_remote_code=self._trust_remote_code,
+            )
+        logger.info(f"Sentence Transformers 模型加载完成: {self._model_name}")
+
+    @property
+    def model(self) -> SentenceTransformer:
+        """获取模型实例（确保已初始化）."""
+        self._ensure_initialized()
+        assert self._model, "模型初始化后仍为 None"
+        return self._model
+
+    def _format_messages(self, messages: list[MessageData]) -> str:
+        """将消息列表格式化为单个文本字符串.
+
+        Args:
+            messages: 消息列表，每个消息包含 role 和 content
+
+        Returns:
+            格式化后的文本字符串
+        """
+        parts: list[str] = []
+        for msg in messages:
+            content = msg.content
+            if isinstance(content, dict):
+                # 如果 content 是字典，尝试提取 text 字段，否则转为字符串
+                content = content.get("text", str(content))
+            parts.append(f"{msg.role}: {content}")
+        return "\n".join(parts)
+
+    def _prepare_texts(self, dataset: list[EmbeddingInputItem]) -> list[str]:
+        """准备输入文本列表.
+
+        Args:
+            dataset: 待嵌入的数据项列表
+
+        Returns:
+            格式化后的文本列表
+        """
+        texts: list[str] = []
+        for item in dataset:
+            text = self._format_messages(item.messages)
+            if self._prompt:
+                text = self._prompt + text
+            texts.append(text)
+        return texts
+
+    @override
+    def embed(self, dataset: list[EmbeddingInputItem]) -> list[EmbeddingResult]:
+        """执行嵌入计算.
+
+        Args:
+            dataset: 待嵌入的数据项列表
+
+        Returns:
+            嵌入结果列表
+        """
+        logger.info(f"开始嵌入计算，数据量: {len(dataset)}")
+
+        # 确保模型已初始化
+        self._ensure_initialized()
+
+        # 准备输入文本
+        logger.debug("准备输入文本...")
+        texts = self._prepare_texts(dataset)
+
+        # 执行嵌入计算
+        logger.info("开始模型推理...")
+        with timing_context("模型推理"):
+            embeddings = self.model.encode(
+                texts,
+                batch_size=self._batch_size,
+                normalize_embeddings=self._normalize_embeddings,
+                show_progress_bar=True,
+            )
+        logger.info(f"嵌入计算完成，输出 {len(embeddings)} 条结果")
+
+        # 构建结果列表
+        return [
+            EmbeddingResult(
+                embedding=embedding.tolist(),
+                data_item=item,
+                meta=item.meta,
+            )
+            for embedding, item in zip(embeddings, dataset)
+        ]

distflow/embed/types.py

@@ -0,0 +1,13 @@
+from typing import Any, TypeAlias
+
+from pydantic import BaseModel
+
+from distflow.data.types import DatasetProcessOutputItem
+
+EmbeddingInputItem: TypeAlias = DatasetProcessOutputItem
+
+
+class EmbeddingResult(BaseModel):  # type: ignore[misc]
+    embedding: list[float]
+    data_item: DatasetProcessOutputItem
+    meta: dict[str, Any]