distflow 0.0.0__tar.gz

distflow-0.0.0/PKG-INFO

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: distflow
+Version: 0.0.0
+Summary: Distance Computation Package for Data Preparation Bench
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: aiohttp
+Requires-Dist: datasets
+Requires-Dist: modelscope
+Requires-Dist: openai
+Requires-Dist: pandas
+Requires-Dist: pydantic
+Requires-Dist: pyyaml
+Requires-Dist: redis
+Requires-Dist: scikit-learn
+Requires-Dist: sentence-transformers
+Requires-Dist: torch
+Requires-Dist: transformers
+Provides-Extra: vllm
+Requires-Dist: vllm; extra == "vllm"
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pyright; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+
+# Data-Preparation-Bench
+
+A benchmark for evaluating the data preparation capabilities of large language models (LLMs). The benchmark is organized into two modules:
+
+## Modules
+
+### 1. Data Synthesis & Augmentation
+
+Given raw metadata, the model is tasked with synthesizing or augmenting datasets to improve downstream model training.
+
+### 2. Data Quality Assessment
+
+Given raw metadata, the model is tasked with predicting the training data's impact on downstream task performance.
+
+## Quick Start
+
+### Usage
+
+The package is published on PyPI and can be installed via pip:
+
+```python
+pip install distflow
+```
+
+For vLLM embedding support, install the optional dependency:
+
+```python
+pip install distflow[vllm]
+```
+
+This project uses [uv](https://docs.astral.sh/uv/) for dependency management. To get started:
+
+```bash
+git clone https://github.com/haolpku/Data-Preparation-Bench.git
+cd Data-Preparation-Bench
+uv sync
+```
+
+To use your own datasets, modify the configuration dictionaries and formatters in [compute_mmd.py](./examples/compute_mmd.py):
+
+```python
+DS1_CONFIG = {
+    "name": "oda-math",
+    "data_path": "OpenDataArena/ODA-Math-460k",
+    "data_size": 5000,
+    "split": "train",
+    "shuffle_seed": 42,
+}
+formatter1 = AlpacaFormatter(
+    user_key="question",
+    assistant_key="response",
+)
+
+DS2_CONFIG = {
+    "name": "infinity-instruct",
+    "data_path": "BAAI/Infinity-Instruct",
+    "data_size": 5000,
+    "split": "train",
+    "shuffle_seed": 42,
+}
+formatter2 = ShareGptFormatter(
+    conversations_key="conversations",
+)
+```
+
+Typically, you only need to update `data_path` with your dataset and define a formatter that converts raw items to the required format. After making these changes, run the MMD computation with:
+
+```bash
+uv run examples/compute_mmd.py
+```
+
+### Development
+
+To set up the development environment locally:
+
+```bash
+uv sync --extra dev
+uv run pre-commit install
+```
+
+Before committing, format and lint the code:
+
+```bash
+uv run pre-commit run --all-files
+```
+
+## Experiment Settings
+
+Please refer to [Experiment.md](./Experiment.md) for detailed experiment configurations.

distflow-0.0.0/README.md

@@ -0,0 +1,89 @@
+# Data-Preparation-Bench
+
+A benchmark for evaluating the data preparation capabilities of large language models (LLMs). The benchmark is organized into two modules:
+
+## Modules
+
+### 1. Data Synthesis & Augmentation
+
+Given raw metadata, the model is tasked with synthesizing or augmenting datasets to improve downstream model training.
+
+### 2. Data Quality Assessment
+
+Given raw metadata, the model is tasked with predicting the training data's impact on downstream task performance.
+
+## Quick Start
+
+### Usage
+
+The package is published on PyPI and can be installed via pip:
+
+```python
+pip install distflow
+```
+
+For vLLM embedding support, install the optional dependency:
+
+```python
+pip install distflow[vllm]
+```
+
+This project uses [uv](https://docs.astral.sh/uv/) for dependency management. To get started:
+
+```bash
+git clone https://github.com/haolpku/Data-Preparation-Bench.git
+cd Data-Preparation-Bench
+uv sync
+```
+
+To use your own datasets, modify the configuration dictionaries and formatters in [compute_mmd.py](./examples/compute_mmd.py):
+
+```python
+DS1_CONFIG = {
+    "name": "oda-math",
+    "data_path": "OpenDataArena/ODA-Math-460k",
+    "data_size": 5000,
+    "split": "train",
+    "shuffle_seed": 42,
+}
+formatter1 = AlpacaFormatter(
+    user_key="question",
+    assistant_key="response",
+)
+
+DS2_CONFIG = {
+    "name": "infinity-instruct",
+    "data_path": "BAAI/Infinity-Instruct",
+    "data_size": 5000,
+    "split": "train",
+    "shuffle_seed": 42,
+}
+formatter2 = ShareGptFormatter(
+    conversations_key="conversations",
+)
+```
+
+Typically, you only need to update `data_path` with your dataset and define a formatter that converts raw items to the required format. After making these changes, run the MMD computation with:
+
+```bash
+uv run examples/compute_mmd.py
+```
+
+### Development
+
+To set up the development environment locally:
+
+```bash
+uv sync --extra dev
+uv run pre-commit install
+```
+
+Before committing, format and lint the code:
+
+```bash
+uv run pre-commit run --all-files
+```
+
+## Experiment Settings
+
+Please refer to [Experiment.md](./Experiment.md) for detailed experiment configurations.

distflow-0.0.0/pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "distflow"
+dynamic = ["version"]
+description = "Distance Computation Package for Data Preparation Bench"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "aiohttp",
+    "datasets",
+    "modelscope",
+    "openai",
+    "pandas",
+    "pydantic",
+    "pyyaml",
+    "redis",
+    "scikit-learn",
+    "sentence-transformers",
+    "torch",
+    "transformers",
+]
+
+[project.optional-dependencies]
+vllm = ["vllm"]
+dev = [
+    "pre-commit",
+    "pyright",
+    "pytest",
+]
+
+[tool.black]
+line-length = 88
+target-version = ['py312']
+include = '\.pyi?$'
+
+[tool.isort]
+profile = "black"
+line_length = 88
+src_paths = ["src", "tests"]
+
+[tool.hatch.version]
+source = "vcs"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/distflow"]

distflow-0.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

distflow-0.0.0/src/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-0.0.0/src/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-0.0.0/src/distflow/data/data_formatter.py

@@ -0,0 +1,112 @@
+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,
+            },
+        )
+
+
+from typing import Any, cast
+
+
+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 not isinstance(conv, dict):
+                continue
+
+            # 检测格式类型并提取字段
+            role = None
+            content = None
+
+            # 标准格式: {"role": "user", "content": "..."}
+            if "role" in conv and "content" in conv:
+                role = conv.get("role")
+                content = conv.get("content")
+
+            # ShareGPT 格式: {"from": "human", "value": "..."}
+            elif "from" in conv and "value" in conv:
+                from_field = conv.get("from")
+                content = conv.get("value")
+
+                assert isinstance(from_field, str) and isinstance(
+                    content, str
+                ), "from和content必须都是str类型"
+
+                role_mapping = {
+                    "human": "user",
+                    "gpt": "assistant",
+                    "system": "system",
+                    "user": "user",
+                    "assistant": "assistant",
+                }
+                role = role_mapping.get(from_field, from_field)
+
+            # 添加到 messages
+            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,
+                "detected_format": (
+                    "sharegpt"
+                    if any(isinstance(c, dict) and "from" in c for c in conversations)
+                    else "standard"
+                ),
+            },
+        )

distflow-0.0.0/src/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-0.0.0/src/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-0.0.0/src/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 | None]:
+        """异步嵌入计算.
+
+        Args:
+            dataset: 待嵌入的数据项列表
+
+        Returns:
+            嵌入结果列表，失败项为 None
+        """