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

"""推理服务 — 支持本地执行和 SSH 远程执行两种模式。"""
import asyncio
import json
from pathlib import Path
from typing import Any

from app.config import get_settings
from app.core.logging import logger

settings = get_settings()


async def generate(
    adapter_path: str,
    prompt: str,
    max_new_tokens: int = 256,
    temperature: float = 0.8,
    top_p: float = 0.95,
    repetition_penalty: float = 1.1,
    do_sample: bool = True,
) -> dict[str, Any]:
    """使用已训练的 adapter 生成文本。"""
    # 从 adapter config 中获取 base model ID
    base_model_id = _get_base_model_id(adapter_path)
    if not base_model_id:
        return {"error": "无法找到基础模型信息，请确保训练任务已完成"}

    if settings.use_remote_compute:
        # 远程执行模式（用 to_thread 避免阻塞事件循环）
        from app.core.remote_executor import run_inference_remote

        adapter_dir = Path(adapter_path)
        adapter_id = adapter_dir.name

        result = await asyncio.to_thread(
            run_inference_remote,
            model_id=base_model_id,
            adapter_id=adapter_id,
            prompt=prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
        )
        if result:
            return result
        return {"error": "Remote inference failed"}

    # 本地执行模式（用 to_thread 避免 GPU 操作阻塞事件循环）
    return await asyncio.to_thread(
        _generate_local,
        adapter_path=adapter_path,
        base_model_id=base_model_id,
        prompt=prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=do_sample,
    )


def _generate_local(
    adapter_path: str,
    base_model_id: str,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    do_sample: bool,
) -> dict[str, Any]:
    """本地执行推理。"""
    try:
        import os
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
        from peft import PeftModel

        tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # CUDA_VISIBLE_DEVICES 由调用方（deploy_service）设置为多卡
        # device_map="auto" 自动将模型层分散到所有可见 GPU
        import torch
        device_map = "auto" if torch.cuda.is_available() else "cpu"
        torch.cuda.set_device(0)

        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            torch_dtype=torch.float16,
            device_map=device_map,
        )
        model = PeftModel.from_pretrained(base_model, adapter_path)
        model.eval()

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        # 构建 stop criteria：遇到角色标记就停止，防止复读
        stop_phrases = ["<|user|>", "<|system|>", "<|assistant|>"]
        stop_token_ids = [tokenizer.encode(p, add_special_tokens=False) for p in stop_phrases]

        class StopOnRoleToken(StoppingCriteria):
            def __call__(self, input_ids, scores, **kwargs):
                gen_seq = input_ids[0].tolist()
                for s_ids in stop_token_ids:
                    if len(gen_seq) >= len(s_ids) and gen_seq[-len(s_ids):] == s_ids:
                        return True
                return False

        stopping_criteria = StoppingCriteriaList([StopOnRoleToken()])

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                do_sample=do_sample,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                stopping_criteria=stopping_criteria,
            )

        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        generated_only = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        # 清理可能残留的角色标记
        for marker in ["<|user|>", "<|system|>", "<|assistant|>"]:
            if marker in generated_only:
                generated_only = generated_only[:generated_only.index(marker)]
        generated_only = generated_only.strip()

        return {
            "prompt": prompt,
            "generated_text": generated_text,
            "generated_only": generated_only,
            "tokens_generated": int(outputs.shape[1] - inputs["input_ids"].shape[1]),
        }

    except Exception as e:
        logger.error(f"Inference failed: {e}")
        return {"error": str(e)}


def _get_base_model_id(adapter_path: str) -> str | None:
    """从 adapter config 中获取 base model ID。"""
    config_path = Path(adapter_path) / "adapter_config.json"
    if config_path.exists():
        with open(config_path) as f:
            cfg = json.load(f)
        return cfg.get("base_model_name_or_path")
    return None


async def get_available_adapters() -> list[dict[str, Any]]:
    """列出所有已训练完成的 adapter（仅显示 status=completed 的任务）。"""
    from app.core.db import async_session, TrainingJobModel
    from sqlalchemy import select

    # 查询数据库中训练完成的任务
    async with async_session() as session:
        result = await session.execute(
            select(TrainingJobModel).where(TrainingJobModel.status == "completed")
        )
        completed_jobs = {job.id: job for job in result.scalars().all()}

    if not completed_jobs:
        return []

    adapters_dir = settings.adapters_dir
    if not adapters_dir.exists():
        return []

    result = []
    for job_id, job in sorted(completed_jobs.items(), key=lambda x: x[1].created_at, reverse=True):
        adapter_dir = adapters_dir / job_id
        if not adapter_dir.is_dir():
            continue
        adapter_config = adapter_dir / "adapter_config.json"
        if not adapter_config.exists():
            continue

        with open(adapter_config) as f:
            cfg = json.load(f)
        result.append({
            "id": job_id,
            "path": str(adapter_dir),
            "base_model": cfg.get("base_model_name_or_path", "unknown"),
            "peft_type": cfg.get("peft_type", "unknown"),
            "model_id": job.model_id,
            "peft_method": job.peft_method,
            "task_type": job.task_type,
            "created_at": job.created_at.isoformat() if job.created_at else None,
        })
    return result


async def run_inference_single(
    model_id: str,
    adapter_id: str,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    do_sample: bool,
) -> dict[str, Any]:
    """供远程 SSH 调用的单条推理入口。"""
    adapter_path = str(settings.adapters_dir / adapter_id)
    return _generate_local(
        adapter_path=adapter_path,
        base_model_id=model_id,
        prompt=prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=do_sample,
    )