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

from pathlib import Path
from typing import Any

from app.core.logging import logger


async def test_model(model_id: str, prompt: str, max_new_tokens: int = 128, temperature: float = 0.8, top_p: float = 0.95) -> dict[str, Any]:
    """加载已缓存模型并生成测试响应。"""
    try:
        import torch
        from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
        from transformers import AutoConfig

        from app.services.model_service import resolve_model_path
        model_path = await resolve_model_path(model_id)
        if not model_path:
            return {"error": f"Model not found in cache: {model_id}"}

        model_dir = Path(model_path)
        if not (model_dir / "config.json").exists():
            return {"error": f"Model directory not found: {model_dir}"}

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

        # 通用加载策略：尝试多种加载方式，自动兼容各种新架构
        model = None
        for loader_cls, kwargs in [
            (AutoModelForCausalLM, {"trust_remote_code": True}),
            (AutoModel, {"trust_remote_code": True}),
        ]:
            try:
                model = loader_cls.from_pretrained(
                    model_dir,
                    torch_dtype=torch.float16,
                    device_map="auto",
                    **kwargs,
                )
                break
            except Exception:
                continue

        if model is None:
            return {"error": f"Unable to load model with any available loader. Model type may not be supported yet."}
        model.eval()

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

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

        generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

        return {
            "model_id": model_id,
            "prompt": prompt,
            "generated_text": generated_text,
        }

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