Fine-tuning/backend/app/engines/text_engine.py

import asyncio
import json
from pathlib import Path
from typing import Any

from app.config import get_settings
from app.core.logging import logger
from app.engines.base import BaseEngine

settings = get_settings()


class TextEngine(BaseEngine):
    """文本模型训练引擎 (LLaMA/Qwen/ChatGLM 等因果语言模型)。"""

    def __init__(self):
        self._tokenizer = None
        self._model = None

    async def load_model(self, model_id: str, **kwargs: Any) -> None:
        """下载并加载基础模型。"""
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        # 优先从数据库获取实际路径（兼容 ModelScope 下载的目录结构）
        from app.services.model_service import resolve_model_path
        model_path = await resolve_model_path(model_id)
        if model_path:
            local_path = model_path
        else:
            local_path = str(settings.models_dir / model_id.replace("/", "_"))

        # 如果本地没有，从 HF 下载
        if not (Path(local_path) / "config.json").exists():
            from huggingface_hub import snapshot_download

            snapshot_download(
                repo_id=model_id,
                local_dir=local_path,
                local_dir_use_symlinks=False,
            )

        quantization = kwargs.get("quantization", None)
        
        # 日志：检查 GPU 状态
        logger.info(f"CUDA available: {torch.cuda.is_available()}")
        logger.info(f"CUDA device count: {torch.cuda.device_count()}")
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")
                logger.info(f"GPU {i} memory: {torch.cuda.get_device_properties(i).total_memory / (1024**3):.2f} GB")
        else:
            logger.warning("No GPU detected! Training will run on CPU.")
        
        max_memory = {i: "4GB" for i in range(torch.cuda.device_count())} if torch.cuda.is_available() else None
        
        load_kwargs: dict[str, Any] = {
            "torch_dtype": torch.float16,
            "device_map": "auto",
            "low_cpu_mem_usage": True,
            "use_safetensors": True,
            "max_memory": max_memory,
        }
        if quantization == "4bit" or quantization == "qlora":
            load_kwargs["load_in_4bit"] = True
            load_kwargs["bnb_4bit_quant_type"] = "nf4"
            load_kwargs["bnb_4bit_use_double_quant"] = True
            load_kwargs["bnb_4bit_compute_dtype"] = torch.float16
        elif quantization == "8bit":
            load_kwargs["load_in_8bit"] = True

        self._tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
        if self._tokenizer.pad_token is None:
            self._tokenizer.pad_token = self._tokenizer.eos_token

        self._model = AutoModelForCausalLM.from_pretrained(local_path, **load_kwargs)
        logger.info(f"Loaded model: {model_id}")

    def get_peft_config(self, method: str, params: dict[str, Any]) -> Any:
        """根据 PEFT 方法返回对应的配置对象。"""
        from app.peft import (
            build_adalora_config,
            build_ia3_config,
            build_lora_config,
            build_prefix_tuning_config,
            build_qlora_config,
        )

        builders = {
            "lora": build_lora_config,
            "qlora": build_qlora_config,
            "ia3": build_ia3_config,
            "adalora": build_adalora_config,
            "prefix_tuning": build_prefix_tuning_config,
        }
        builder = builders.get(method, build_lora_config)
        return builder(params)

    async def preprocess_dataset(
        self,
        dataset_path: str,
        output_path: str,
        task_type: str = "sft",
        template: str = "alpaca",
        **kwargs: Any,
    ) -> str:
        """将数据集预处理为训练格式。"""
        from app.preprocessors import preprocess_file

        processed = preprocess_file(dataset_path, output_path, task_type, template)
        logger.info(f"Preprocessed {len(processed)} samples for {task_type}/{template}")
        return output_path

    async def train(
        self,
        job_id: str,
        dataset_path: str,
        peft_config: Any,
        training_args: dict[str, Any],
    ) -> str:
        """执行训练。"""
        from peft import get_peft_model
        from transformers import DataCollatorForSeq2Seq, TrainingArguments

        task_type = training_args.get("task_type", "sft")
        epochs = training_args.get("epochs", 3)
        batch_size = training_args.get("batch_size", 4)
        gradient_accumulation = training_args.get("gradient_accumulation", 4)
        learning_rate = training_args.get("learning_rate", 2e-4)
        max_seq_length = training_args.get("max_seq_length", 2048)
        warmup_ratio = training_args.get("warmup_ratio", 0.05)
        save_strategy = training_args.get("save_strategy", "epoch")
        deepspeed_config = training_args.get("deepspeed", None)

        dataset = self._tokenize_dataset(dataset_path, max_seq_length)

        self._model = get_peft_model(self._model, peft_config)
        self._model.print_trainable_parameters()

        output_dir = str(settings.adapters_dir / job_id)
        tr_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=epochs,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=gradient_accumulation,
            learning_rate=learning_rate,
            warmup_ratio=warmup_ratio,
            save_strategy=save_strategy,
            logging_strategy="steps",
            logging_steps=10,
            fp16=True,
            optim="adamw_torch",
            remove_unused_columns=False,
            report_to="none",
            gradient_checkpointing=True,
            gradient_checkpointing_kwargs={"use_reentrant": False},
            dataloader_num_workers=1,
            dataloader_pin_memory=False,
            **({"deepspeed": deepspeed_config} if deepspeed_config else {}),
        )

        callback = _ProgressCallback(job_id)

        if task_type == "sft":
            from transformers import Trainer

            trainer = Trainer(
                model=self._model,
                args=tr_args,
                train_dataset=dataset,
                data_collator=DataCollatorForSeq2Seq(self._tokenizer),
                callbacks=[callback],
            )
        else:
            from trl import (
                DPOConfig,
                DPOTrainer,
                KTOConfig,
                KTOTrainer,
                ORPOConfig,
                ORPOTrainer,
            )

            base_trainer_kwargs = dict(
                output_dir=output_dir,
                num_train_epochs=epochs,
                per_device_train_batch_size=batch_size,
                gradient_accumulation_steps=gradient_accumulation,
                learning_rate=learning_rate,
                warmup_ratio=warmup_ratio,
                save_strategy=save_strategy,
                logging_steps=10,
                fp16=True,
                report_to="none",
            )

            if task_type == "dpo":
                trainer = DPOTrainer(
                    model=self._model,
                    args=DPOConfig(**base_trainer_kwargs),
                    train_dataset=dataset,
                    processing_class=self._tokenizer,
                )
            elif task_type == "orpo":
                trainer = ORPOTrainer(
                    model=self._model,
                    args=ORPOConfig(**base_trainer_kwargs),
                    train_dataset=dataset,
                    processing_class=self._tokenizer,
                )
            elif task_type == "kto":
                trainer = KTOTrainer(
                    model=self._model,
                    args=KTOConfig(**base_trainer_kwargs),
                    train_dataset=dataset,
                    processing_class=self._tokenizer,
                )
            else:
                trainer = Trainer(
                    model=self._model,
                    args=tr_args,
                    train_dataset=dataset,
                    data_collator=DataCollatorForSeq2Seq(self._tokenizer),
                    callbacks=[callback],
                )

        try:
            trainer.train()
            self._model.save_pretrained(output_dir)
            self._tokenizer.save_pretrained(output_dir)
            logger.info(f"Training completed for job {job_id}")
        except Exception as e:
            logger.error(f"Training failed for job {job_id}: {e}")
            raise

        return output_dir

    def get_model_info(self, model_id: str) -> dict[str, Any]:
        """读取模型配置信息。"""
        import json
        from pathlib import Path

        # 同步查找模型路径（兼容 HF 和 ModelScope）
        candidates = [
            settings.models_dir / model_id.replace("/", "_"),
            settings.models_dir / model_id,
        ]
        config_path = None
        for p in candidates:
            if (p / "config.json").exists():
                config_path = p / "config.json"
                break
        if not config_path:
            # 最后尝试扫描
            model_name = model_id.split("/")[-1]
            for cp in settings.models_dir.rglob("config.json"):
                if model_name in str(cp.parent):
                    config_path = cp
                    break

        if config_path.exists():
            with open(config_path) as f:
                config = json.load(f)
            return {
                "model_type": config.get("model_type", "causal_lm"),
                "context_length": config.get("max_position_embeddings", config.get("max_sequence_length", 2048)),
                "hidden_size": config.get("hidden_size", 0),
                "num_layers": config.get("num_hidden_layers", 0),
            }
        return {"model_type": "causal_lm", "context_length": 2048}

    def _tokenize_dataset(self, dataset_path: str, max_seq_length: int):
        """Tokenize 处理后的 JSONL 数据集。"""
        from datasets import Dataset as HFDataset

        data = []
        with open(dataset_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    item = json.loads(line)
                    # 确保 prompt 和 completion 是字符串
                    if "prompt" in item:
                        if isinstance(item["prompt"], (list, dict)):
                            item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False)
                        item["prompt"] = str(item["prompt"])
                    if "completion" in item:
                        if isinstance(item["completion"], (list, dict)):
                            item["completion"] = json.dumps(item["completion"], ensure_ascii=False)
                        item["completion"] = str(item["completion"])
                    data.append(item)

        hf_dataset = HFDataset.from_list(data)

        def tokenize_fn(batch):
            def _to_str(v):
                if isinstance(v, (list, dict)):
                    return json.dumps(v, ensure_ascii=False)
                return str(v) if v is not None else ""

            raw_prompts = batch.get("prompt", [])
            raw_completions = batch.get("completion", [])

            prompts = [_to_str(v) for v in raw_prompts]
            completions = [_to_str(v) for v in raw_completions]

            if not prompts:
                return {"input_ids": [], "attention_mask": [], "labels": []}

            full_texts = [f"{p}\n{c}" for p, c in zip(prompts, completions)]
            tokenized = self._tokenizer(
                full_texts, truncation=True, max_length=max_seq_length, padding=False,
            )
            tokenized["labels"] = list(tokenized["input_ids"])
            return tokenized

        tokenized_dataset = hf_dataset.map(
            tokenize_fn,
            batched=True,
            remove_columns=["prompt", "completion"],
        )
        return tokenized_dataset


class _ProgressCallback:
    """自定义训练进度回调，通过 WebSocket 发送进度。"""

    def __init__(self, job_id: str):
        self.job_id = job_id

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs and "loss" in logs:
            asyncio.create_task(
                send_progress(
                    self.job_id,
                    epoch=int(state.epoch or 0),
                    step=state.global_step,
                    total_steps=state.max_steps or 0,
                    loss=logs["loss"],
                    learning_rate=logs.get("learning_rate", 0),
                )
            )

    def on_epoch_end(self, args, state, control, **kwargs):
        asyncio.create_task(
            send_epoch_done(self.job_id, epoch=int(state.epoch or 0), eval_loss=None, eval_accuracy=None)
        )

    def on_train_end(self, args, state, control, **kwargs):
        asyncio.create_task(
            send_completed(
                self.job_id,
                total_time_seconds=getattr(state, "train_runtime", 0),
                adapter_path=str(settings.adapters_dir / self.job_id),
            )
        )

    def on_train_begin(self, args, state, control, **kwargs):
        pass

    def on_step_begin(self, args, state, control, **kwargs):
        pass

    def on_step_end(self, args, state, control, **kwargs):
        pass

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        pass

    def on_save(self, args, state, control, **kwargs):
        pass

    def on_predict(self, args, state, control, metrics=None, **kwargs):
        pass

    def on_init_end(self, args, state, control, **kwargs):
        pass

    def on_epoch_begin(self, args, state, control, **kwargs):
        pass


# 全局单例
text_engine = TextEngine()