275 lines
9.0 KiB
Python
275 lines
9.0 KiB
Python
"""
|
|
数据导出服务 - 负责将合成数据导出为各种格式
|
|
"""
|
|
import datetime
|
|
import json
|
|
import os
|
|
from typing import Iterable, List, Optional, Sequence
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.logging import get_logger
|
|
from app.db.models.data_synthesis import (
|
|
DataSynthInstance,
|
|
DataSynthesisFileInstance,
|
|
SynthesisData,
|
|
)
|
|
from app.db.models.dataset_management import Dataset, DatasetFiles
|
|
from app.module.generation.schema.generation import (
|
|
ExportSynthesisDataResponse,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class SynthesisExportError(Exception):
|
|
"""导出失败时抛出的异常"""
|
|
|
|
|
|
class SynthesisDatasetExporter:
|
|
"""
|
|
数据导出器 - 将合成数据导出为各种格式
|
|
|
|
导出规则：
|
|
- 维度：原始文件 (DatasetFiles)
|
|
- 每个原始文件生成一个 JSONL 文件
|
|
- JSONL 文件名称与原始文件名称一致
|
|
- 支持格式转换：alpaca、sharegpt、raw
|
|
"""
|
|
|
|
SUPPORTED_FORMATS = ["alpaca", "sharegpt", "raw"]
|
|
DEFAULT_FORMAT = "alpaca"
|
|
|
|
def __init__(self, db: AsyncSession, format: str = "alpaca", output_path: Optional[str] = None):
|
|
self._db = db
|
|
self._format = format if format in self.SUPPORTED_FORMATS else self.DEFAULT_FORMAT
|
|
self._output_path = output_path
|
|
|
|
async def export_task_to_dataset(
|
|
self,
|
|
task_id: str,
|
|
dataset_id: str,
|
|
) -> Dataset:
|
|
"""
|
|
将任务的全部合成数据导出到已有数据集
|
|
|
|
Args:
|
|
task_id: 任务 ID
|
|
dataset_id: 目标数据集 ID
|
|
|
|
Returns:
|
|
更新后的数据集
|
|
"""
|
|
task = await self._db.get(DataSynthInstance, task_id)
|
|
if not task:
|
|
raise SynthesisExportError(f"Synthesis task {task_id} not found")
|
|
|
|
dataset = await self._db.get(Dataset, dataset_id)
|
|
if not dataset:
|
|
raise SynthesisExportError(f"Dataset {dataset_id} not found")
|
|
|
|
file_instances = await self._load_file_instances(task_id)
|
|
if not file_instances:
|
|
raise SynthesisExportError("No synthesis file instances found for task")
|
|
|
|
base_path = self._ensure_dataset_path(dataset)
|
|
|
|
created_files: list[DatasetFiles] = []
|
|
total_size = 0
|
|
|
|
for file_instance in file_instances:
|
|
records = await self._load_synthesis_data_for_file(file_instance.id)
|
|
if not records:
|
|
continue
|
|
|
|
original_name = file_instance.file_name or "unknown"
|
|
base_name, _ = os.path.splitext(original_name)
|
|
archived_file_name = f"{base_name}.jsonl"
|
|
|
|
file_path = os.path.join(base_path, archived_file_name)
|
|
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
|
self._write_jsonl(file_path, records, self._format)
|
|
|
|
try:
|
|
file_size = os.path.getsize(file_path)
|
|
except OSError:
|
|
file_size = 0
|
|
|
|
df = DatasetFiles(
|
|
dataset_id=dataset.id,
|
|
file_name=archived_file_name,
|
|
file_path=file_path,
|
|
file_type="jsonl",
|
|
file_size=file_size,
|
|
last_access_time=datetime.datetime.now(),
|
|
)
|
|
self._db.add(df)
|
|
created_files.append(df)
|
|
total_size += file_size
|
|
|
|
if created_files:
|
|
dataset.file_count = (dataset.file_count or 0) + len(created_files)
|
|
dataset.size_bytes = (dataset.size_bytes or 0) + total_size
|
|
dataset.status = "ACTIVE"
|
|
|
|
await self._db.commit()
|
|
|
|
logger.info(
|
|
"Exported synthesis task %s to dataset %s with %d files (total %d bytes)",
|
|
task_id,
|
|
dataset.id,
|
|
len(created_files),
|
|
total_size,
|
|
)
|
|
|
|
return dataset
|
|
|
|
async def export_data(
|
|
self,
|
|
task_id: str,
|
|
file_instance_ids: Optional[List[str]] = None,
|
|
) -> ExportSynthesisDataResponse:
|
|
"""
|
|
将合成数据导出为指定格式的 JSONL 文件
|
|
|
|
Args:
|
|
task_id: 任务 ID
|
|
file_instance_ids: 文件实例 ID 列表，为空则导出全部
|
|
|
|
Returns:
|
|
导出结果
|
|
"""
|
|
if file_instance_ids:
|
|
file_instances = await self._load_specific_file_instances(file_instance_ids)
|
|
else:
|
|
file_instances = await self._load_file_instances(task_id)
|
|
|
|
if not file_instances:
|
|
raise SynthesisExportError("No file instances to export")
|
|
|
|
output_dir = self._output_path or self._get_default_output_dir()
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
file_paths: List[str] = []
|
|
total_records = 0
|
|
|
|
for file_instance in file_instances:
|
|
records = await self._load_synthesis_data_for_file(file_instance.id)
|
|
if not records:
|
|
continue
|
|
|
|
original_name = file_instance.file_name or "unknown"
|
|
base_name, _ = os.path.splitext(original_name)
|
|
output_file_name = f"{base_name}.jsonl"
|
|
|
|
file_path = os.path.join(output_dir, output_file_name)
|
|
self._write_jsonl(file_path, records, self._format)
|
|
|
|
file_paths.append(file_path)
|
|
total_records += len(records)
|
|
|
|
return ExportSynthesisDataResponse(
|
|
file_paths=file_paths,
|
|
total_records=total_records,
|
|
format=self._format,
|
|
)
|
|
|
|
async def _load_file_instances(self, task_id: str) -> Sequence[DataSynthesisFileInstance]:
|
|
"""加载任务下的所有文件实例"""
|
|
result = await self._db.execute(
|
|
select(DataSynthesisFileInstance).where(
|
|
DataSynthesisFileInstance.synthesis_instance_id == task_id
|
|
)
|
|
)
|
|
return result.scalars().all()
|
|
|
|
async def _load_specific_file_instances(
|
|
self, file_ids: List[str]
|
|
) -> Sequence[DataSynthesisFileInstance]:
|
|
"""加载指定的文件实例"""
|
|
result = await self._db.execute(
|
|
select(DataSynthesisFileInstance).where(
|
|
DataSynthesisFileInstance.id.in_(file_ids)
|
|
)
|
|
)
|
|
return result.scalars().all()
|
|
|
|
async def _load_synthesis_data_for_file(
|
|
self, file_instance_id: str
|
|
) -> List[dict]:
|
|
"""加载单个文件实例的所有合成数据"""
|
|
result = await self._db.execute(
|
|
select(SynthesisData).where(
|
|
SynthesisData.synthesis_file_instance_id == file_instance_id
|
|
)
|
|
)
|
|
rows: Sequence[SynthesisData] = result.scalars().all()
|
|
|
|
records: List[dict] = []
|
|
for row in rows:
|
|
payload = row.data or {}
|
|
records.append(payload)
|
|
return records
|
|
|
|
def _write_jsonl(self, path: str, records: Iterable[dict], format: Optional[str] = None) -> None:
|
|
"""写入 JSONL 文件"""
|
|
fmt = format or self._format
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
for record in records:
|
|
formatted = self._format_record(record, fmt)
|
|
f.write(json.dumps(formatted, ensure_ascii=False))
|
|
f.write("\n")
|
|
|
|
def _format_record(self, record: dict, format: str) -> dict:
|
|
"""根据格式转换记录"""
|
|
if format == "alpaca":
|
|
return self._format_as_alpaca(record)
|
|
elif format == "sharegpt":
|
|
return self._format_as_sharegpt(record)
|
|
else:
|
|
return self._format_as_raw(record)
|
|
|
|
def _format_as_alpaca(self, record: dict) -> dict:
|
|
"""转换为 Alpaca 格式"""
|
|
return {
|
|
"instruction": record.get("instruction", ""),
|
|
"output": record.get("output", ""),
|
|
}
|
|
|
|
def _format_as_sharegpt(self, record: dict) -> dict:
|
|
"""转换为 ShareGPT 格式"""
|
|
instruction = record.get("instruction", "")
|
|
output = record.get("output", "")
|
|
input_text = record.get("input", "")
|
|
|
|
conversations = [
|
|
{"from": "human", "value": instruction},
|
|
{"from": "gpt", "value": output},
|
|
]
|
|
|
|
result = {"conversations": conversations}
|
|
|
|
if input_text:
|
|
result["context"] = input_text
|
|
|
|
return result
|
|
|
|
def _format_as_raw(self, record: dict) -> dict:
|
|
"""原始格式"""
|
|
return record
|
|
|
|
def _get_default_output_dir(self) -> str:
|
|
"""获取默认输出目录"""
|
|
import tempfile
|
|
return tempfile.gettempdir()
|
|
|
|
@staticmethod
|
|
def _ensure_dataset_path(dataset: Dataset) -> str:
|
|
"""确保数据集路径存在"""
|
|
if not dataset.path:
|
|
raise SynthesisExportError("Dataset path is empty")
|
|
os.makedirs(dataset.path, exist_ok=True)
|
|
return dataset.path
|