Fine-tuning/backend/app/services/dataset_service.py

import asyncio
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


from fastapi import UploadFile

from app.config import get_settings
from app.core.db import async_session, DatasetRecord
from app.core.logging import logger
from app.schemas.dataset import DatasetDownloadRequest, DatasetDownloadResponse

settings = get_settings()


async def download_dataset(req: DatasetDownloadRequest) -> DatasetDownloadResponse:
    """从 HuggingFace 或 ModelScope 下载数据集。"""
    try:
        if req.use_modelscope:
            # 用 asyncio.to_thread 包裹同步下载，避免阻塞事件循环
            ds_dir, jsonl_path, record_count = await asyncio.to_thread(_download_modelscope_dataset, req.dataset_id)
        else:
            from datasets import load_dataset

            ds = load_dataset(req.dataset_id)
            ds_dir = settings.processed_dir / f"hf_{req.dataset_id.replace('/', '_')}"
            ds_dir.mkdir(parents=True, exist_ok=True)
            if "train" in ds:
                split = ds["train"]
            else:
                split = ds[list(ds.keys())[0]]
            output_path = ds_dir / "data.jsonl"
            with open(output_path, "w", encoding="utf-8") as f:
                for item in split:
                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
            jsonl_path = output_path
            record_count = len(split) if hasattr(split, "__len__") else 0

        # 写入数据库
        record = DatasetRecord(
            id=str(uuid.uuid4()),
            name=req.dataset_id,
            format="jsonl",
            record_count=record_count,
            file_path=str(jsonl_path),
            created_at=datetime.now(timezone.utc),
        )
        async with async_session() as session:
            session.add(record)
            await session.commit()

        logger.info(f"Downloaded dataset: {req.dataset_id} ({record_count} records, source={'ModelScope' if req.use_modelscope else 'HuggingFace'})")
        return DatasetDownloadResponse(dataset_id=req.dataset_id, status="completed", path=str(jsonl_path))
    except Exception as e:
        logger.error(f"Dataset download failed: {e}")
        return DatasetDownloadResponse(dataset_id=req.dataset_id, status="failed", error=str(e))


def _download_modelscope_dataset(dataset_id: str) -> tuple[Path, Path, int]:
    """用 snapshot_download 下载数据集文件，完全绕过 datasets 库，避免版本兼容问题。"""
    from modelscope import snapshot_download

    ds_dir = settings.processed_dir / f"ms_{dataset_id.replace('/', '_')}"
    ds_dir.mkdir(parents=True, exist_ok=True)

    # 用 snapshot_download 下载数据集文件到本地
    local_path = snapshot_download(dataset_id, cache_dir=str(settings.processed_dir))

    # 扫描下载目录中的 JSON/JSONL 文件
    data_files = []
    for p in Path(local_path).rglob("*"):
        if p.is_file() and p.suffix in (".json", ".jsonl"):
            data_files.append(p)

    if not data_files:
        raise ValueError(f"No JSON/JSONL files found in dataset {dataset_id}")

    # 优先取 train / data 开头的文件
    target = None
    for name in ("train.jsonl", "train.json", "data.jsonl", "data.json"):
        for f in data_files:
            if f.name == name:
                target = f
                break
        if target:
            break
    if not target:
        target = data_files[0]

    # 读取并统一转为 JSONL
    jsonl_path = ds_dir / "data.jsonl"
    record_count = 0
    content = target.read_text(encoding="utf-8")
    if target.suffix == ".jsonl":
        records = [json.loads(line.strip()) for line in content.splitlines() if line.strip()]
    else:
        records = json.loads(content)
        if not isinstance(records, list):
            records = [records]

    with open(jsonl_path, "w", encoding="utf-8") as f:
        for item in records:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
            record_count += 1

    return ds_dir, jsonl_path, record_count


async def upload_dataset(file: UploadFile) -> dict[str, Any]:
    """保存上传文件并写入数据库。"""
    upload_dir = settings.uploads_dir
    upload_dir.mkdir(parents=True, exist_ok=True)

    # 避免文件名冲突
    safe_name = file.filename or "unknown"
    file_path = upload_dir / safe_name
    if file_path.exists():
        file_path = upload_dir / f"{uuid.uuid4().hex}_{safe_name}"

    content = await file.read()
    file_path.write_bytes(content)

    fmt = _detect_format(safe_name)
    record_count = _count_records(file_path, fmt)

    record_id = str(uuid.uuid4())
    record = DatasetRecord(
        id=record_id,
        name=safe_name,
        format=fmt,
        record_count=record_count,
        file_path=str(file_path),
        created_at=datetime.now(timezone.utc),
    )
    async with async_session() as session:
        session.add(record)
        await session.commit()

    logger.info(f"Uploaded dataset: {safe_name} ({record_count} records, format={fmt})")
    return {
        "id": record_id,
        "name": safe_name,
        "format": fmt,
        "record_count": record_count,
        "file_path": str(file_path),
        "created_at": record.created_at.isoformat(),
    }


def _format_value(value) -> str:
    """将复杂值格式化为可读字符串。"""
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False, indent=2)
    return str(value)


def _is_sharegpt_format(records: list[dict]) -> bool:
    """检测是否为 ShareGPT 格式。"""
    if not records:
        return False
    first = records[0]
    if "conversations" in first and isinstance(first["conversations"], list):
        if len(first["conversations"]) > 0 and isinstance(first["conversations"][0], dict):
            conv = first["conversations"][0]
            return "from" in conv and "value" in conv
    return False


def _flatten_sharegpt(records: list[dict]) -> tuple[list[dict], list[str]]:
    """将 ShareGPT 格式展平为 input/output 列。"""
    flat_rows = []
    for row in records:
        conversations = row.get("conversations", [])
        # 每轮 user+assistant 对话作为一行
        for i in range(0, len(conversations) - 1, 2):
            user_turn = conversations[i]
            assistant_turn = conversations[i + 1] if i + 1 < len(conversations) else None

            if user_turn.get("from") in ("human", "user"):
                input_text = str(user_turn.get("value", ""))
                output_text = str(assistant_turn.get("value", "")) if assistant_turn else ""
            else:
                input_text = str(assistant_turn.get("value", "")) if assistant_turn else ""
                output_text = str(user_turn.get("value", ""))

            # 截断过长文本
            if len(input_text) > 500:
                input_text = input_text[:500] + "..."
            if len(output_text) > 500:
                output_text = output_text[:500] + "..."

            flat_rows.append({"input": input_text, "output": output_text})

    return flat_rows, ["input", "output"]


async def preview_dataset(dataset_id: str, rows: int = 10) -> dict[str, Any]:
    """预览数据集前 N 行。"""
    async with async_session() as session:
        from sqlalchemy import select

        result = await session.execute(select(DatasetRecord).where(DatasetRecord.id == dataset_id))
        record = result.scalar_one_or_none()
        if not record:
            return {"total_records": 0, "preview_rows": [], "columns": []}

    file_path = Path(record.file_path)
    if not file_path.exists():
        return {"total_records": 0, "preview_rows": [], "columns": []}

    fmt = record.format
    preview_data = _read_records(file_path, fmt, rows)

    # 检测是否为 ShareGPT 格式，如果是则展平为 input/output 列
    if _is_sharegpt_format(preview_data):
        preview_data, columns = _flatten_sharegpt(preview_data)
    else:
        columns = list(preview_data[0].keys()) if preview_data else []

    return {
        "total_records": record.record_count,
        "preview_rows": [
            {
                "row_index": i,
                "data": {k: _format_value(v) for k, v in row.items()},
            }
            for i, row in enumerate(preview_data)
        ],
        "columns": columns,
    }


async def validate_dataset(dataset_id: str) -> dict[str, Any]:
    """校验数据集格式和 Schema。"""
    async with async_session() as session:
        from sqlalchemy import select

        result = await session.execute(select(DatasetRecord).where(DatasetRecord.id == dataset_id))
        record = result.scalar_one_or_none()
        if not record:
            return {"is_valid": False, "errors": ["Dataset not found"], "warnings": []}

    file_path = Path(record.file_path)
    if not file_path.exists():
        return {"is_valid": False, "errors": ["File not found"], "warnings": []}

    errors = []
    warnings = []

    # 检查格式
    fmt = record.format
    if fmt not in ("jsonl", "csv", "json", "parquet"):
        errors.append(f"Unsupported format: {fmt}")

    # 检查内容
    try:
        preview = _read_records(file_path, fmt, 5)
        if not preview:
            warnings.append("Dataset appears to be empty")
        else:
            # 检查必需字段（SFT 格式）
            first = preview[0]
            has_sft_fields = any(k in first for k in ("instruction", "prompt", "text", "input", "output", "completion"))
            if not has_sft_fields:
                warnings.append(f"No common SFT fields found. Keys: {list(first.keys())}")
    except Exception as e:
        errors.append(f"Failed to read file: {str(e)}")

    return {"is_valid": len(errors) == 0, "errors": errors, "warnings": warnings}


async def list_datasets() -> list[dict[str, Any]]:
    """列出所有已上传数据集。"""
    async with async_session() as session:
        from sqlalchemy import select

        result = await session.execute(select(DatasetRecord).order_by(DatasetRecord.created_at.desc()))
        records = result.scalars().all()

    return [
        {
            "id": r.id,
            "name": r.name,
            "format": r.format,
            "record_count": r.record_count,
            "file_path": r.file_path,
            "created_at": r.created_at.isoformat(),
        }
        for r in records
    ]


async def delete_dataset(dataset_id: str) -> dict[str, Any]:
    """删除数据集。"""
    async with async_session() as session:
        from sqlalchemy import select

        result = await session.execute(select(DatasetRecord).where(DatasetRecord.id == dataset_id))
        record = result.scalar_one_or_none()
        if record:
            # 删除文件
            file_path = Path(record.file_path)
            if file_path.exists():
                file_path.unlink()
            await session.delete(record)
            await session.commit()
            logger.info(f"Deleted dataset: {record.name}")

    return {"status": "deleted"}


def _detect_format(filename: str) -> str:
    ext = Path(filename).suffix.lower().lstrip(".")
    if ext in ("jsonl", "csv", "parquet", "json"):
        return ext
    return "unknown"


def _count_records(file_path: Path, fmt: str) -> int:
    try:
        if fmt == "jsonl":
            return sum(1 for line in open(file_path, encoding="utf-8") if line.strip())
        elif fmt == "json":
            with open(file_path, encoding="utf-8") as f:
                data = json.load(f)
                return len(data) if isinstance(data, list) else 1
        elif fmt == "csv":
            import csv
            with open(file_path, encoding="utf-8") as f:
                return sum(1 for _ in csv.reader(f)) - 1  # minus header
        elif fmt == "parquet":
            import pandas as pd
            return len(pd.read_parquet(file_path))
    except Exception:
        pass
    return 0


def _read_records(file_path: Path, fmt: str, n: int) -> list[dict]:
    if fmt == "jsonl":
        records = []
        with open(file_path, encoding="utf-8") as f:
            for i, line in enumerate(f):
                if i >= n:
                    break
                line = line.strip()
                if line:
                    records.append(json.loads(line))
        return records
    elif fmt == "json":
        with open(file_path, encoding="utf-8") as f:
            data = json.load(f)
            return data[:n] if isinstance(data, list) else [data]
    elif fmt == "csv":
        import csv
        with open(file_path, encoding="utf-8") as f:
            reader = csv.DictReader(f)
            return [dict(row) for i, row in enumerate(reader) if i < n]
    elif fmt == "parquet":
        import pandas as pd
        df = pd.read_parquet(file_path)
        return df.head(n).to_dict(orient="records")
    return []