core/local_inference.py

#!/usr/bin/env python3
"""
虫群本地推理后端
- 加载训练好的SwarmModel模型
- 提供统一的推理接口
- 支持多模型管理和切换
- 与MOA引擎集成
"""

import os
import sys
import time
import json
import torch
import threading
from typing import Dict, List, Optional, Tuple

sys.path.insert(0, "/home/admin/swarm")

from training.model import SwarmModel
from training.tokenizer import SwarmTokenizer


class LocalModelBackend:
    """
    本地模型推理后端
    
    管理多个本地小模型的加载、推理和卸载。
    按需加载模型到内存，支持LRU淘汰。
    """

    def __init__(self, models_dir: str = "/home/admin/swarm/training/models", max_loaded: int = 3):
        self.models_dir = models_dir
        self.max_loaded = max_loaded
        
        # 已加载的模型: model_id -> (model, tokenizer)
        self._loaded: Dict[str, Tuple[SwarmModel, SwarmTokenizer]] = {}
        self._load_times: Dict[str, float] = {}  # LRU时间戳
        self._lock = threading.Lock()
        
        # 可用模型列表
        self._available = self._scan_models()
        
        # 设备
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def _scan_models(self) -> Dict[str, dict]:
        """扫描可用模型"""
        available = {}
        if not os.path.exists(self.models_dir):
            return available
            
        # 共享分词器路径
        shared_tok = os.path.join(self.models_dir, "shared_tokenizer.json")
        # v2分词器
        v2_tok = os.path.join(self.models_dir, "tokenizer_v2.json")
            
        for name in os.listdir(self.models_dir):
            model_dir = os.path.join(self.models_dir, name)
            if not os.path.isdir(model_dir):
                continue
            model_path = os.path.join(model_dir, "model.pt")
            meta_path = os.path.join(model_dir, "meta.json")
            
            if os.path.exists(model_path):
                meta = {}
                if os.path.exists(meta_path):
                    with open(meta_path, "r") as f:
                        meta = json.load(f)
                # 优先用模型目录下的分词器，否则用共享分词器或v2分词器
                tok_path = os.path.join(model_dir, "tokenizer.json")
                if not os.path.exists(tok_path) and os.path.exists(shared_tok):
                    tok_path = shared_tok
                if not os.path.exists(tok_path) and os.path.exists(v2_tok):
                    tok_path = v2_tok
                # 检测量化模型
                quantized = os.path.exists(os.path.join(model_dir, "quantized_int8"))
                available[name] = {
                    "path": model_dir,
                    "model_path": model_path,
                    "tokenizer_path": tok_path,
                    "meta": meta,
                    "quantized": quantized,
                }
        
        return available

    def list_available(self) -> List[str]:
        """列出可用模型"""
        return list(self._available.keys())

    def get_model_info(self, model_id: str) -> Optional[dict]:
        """获取模型信息"""
        if model_id in self._available:
            return self._available[model_id].get("meta", {})
        if model_id in self._loaded:
            model, _ = self._loaded[model_id]
            return model.get_info()
        return None

    def load(self, model_id: str) -> bool:
        """加载模型到内存"""
        if model_id in self._loaded:
            self._load_times[model_id] = time.time()
            return True
            
        if model_id not in self._available:
            print(f"  ❌ 模型不存在: {model_id}")
            return False
        
        # LRU淘汰
        with self._lock:
            while len(self._loaded) >= self.max_loaded:
                oldest = min(self._load_times, key=self._load_times.get)
                self.unload(oldest)
        
        try:
            info = self._available[model_id]
            
            # 加载分词器
            tokenizer = SwarmTokenizer.load(info["tokenizer_path"])
            
            # 加载模型
            checkpoint = torch.load(info["model_path"], map_location=self.device, weights_only=False)
            saved_cfg = checkpoint.get("config", {})
            
            # 从权重推断实际max_len（保存的max_len可能不准）
            actual_max_len = 256
            state_dict = checkpoint.get("model_state_dict", {})
            if "pos_emb.weight" in state_dict:
                actual_max_len = state_dict["pos_emb.weight"].shape[0]
            elif "max_len" in checkpoint:
                actual_max_len = checkpoint["max_len"]
            
            model = SwarmModel(
                vocab_size=checkpoint.get("vocab_size", tokenizer.vocab_size_actual),
                d_model=saved_cfg.get("d_model", 192 if model_id == "tiny" else 320),
                n_heads=saved_cfg.get("n_heads", 6 if model_id == "tiny" else 8),
                n_layers=saved_cfg.get("n_layers", 6 if model_id == "tiny" else 8),
                d_ff=saved_cfg.get("d_model", 192) * 4,  # 从d_model推算
                max_len=actual_max_len,
            )
            model.load_state_dict(checkpoint["model_state_dict"])
            model = model.to(self.device)
            model.eval()
            
            with self._lock:
                self._loaded[model_id] = (model, tokenizer)
                self._load_times[model_id] = time.time()
            
            param_count = model.get_info()["params_M"]
            print(f"  ✅ 模型已加载: {model_id} ({param_count}M参数)")
            return True
            
        except Exception as e:
            print(f"  ❌ 加载失败: {model_id}, {e}")
            return False

    def unload(self, model_id: str) -> bool:
        """卸载模型"""
        with self._lock:
            if model_id in self._loaded:
                del self._loaded[model_id]
                del self._load_times[model_id]
                # 释放GPU内存
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                print(f"  🗑️ 模型已卸载: {model_id}")
                return True
        return False

    def infer(
        self,
        model_id: str,
        query: str,
        max_new_tokens: int = 128,
        temperature: float = 0.7,
        top_k: int = 40,
    ) -> Dict:
        """
        推理接口
        
        返回:
        {
            "response": str,
            "model_id": str,
            "latency_ms": float,
            "tokens_generated": int,
            "success": bool,
        }
        """
        # 确保模型已加载
        if model_id not in self._loaded:
            if not self.load(model_id):
                return {
                    "response": "",
                    "model_id": model_id,
                    "latency_ms": 0,
                    "tokens_generated": 0,
                    "success": False,
                }
        
        model, tokenizer = self._loaded[model_id]
        
        start_time = time.perf_counter()
        
        try:
            # 编码输入（只加BOS，不加EOS——EOS是生成终止标记）
            bos_id = tokenizer.bos_id
            input_ids = [bos_id] + tokenizer.encode(query, add_special=False)
            input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)
            
            # 生成（启用重复惩罚，小模型必备）
            with torch.no_grad():
                output_ids = model.generate(
                    input_tensor,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k,
                    eos_id=tokenizer.eos_id,
                    repetition_penalty=1.3,
                )
            
            # 解码
            output_ids = output_ids[0].cpu().tolist()
            # 去掉输入部分
            generated_ids = output_ids[len(input_ids):]
            response = tokenizer.decode(generated_ids)
            
            # 清理
            response = response.replace("[UNK]", "").strip()
            
            latency_ms = (time.perf_counter() - start_time) * 1000
            
            return {
                "response": response,
                "model_id": model_id,
                "latency_ms": round(latency_ms, 1),
                "tokens_generated": len(generated_ids),
                "success": True,
            }
            
        except Exception as e:
            latency_ms = (time.perf_counter() - start_time) * 1000
            return {
                "response": f"推理错误: {str(e)[:100]}",
                "model_id": model_id,
                "latency_ms": round(latency_ms, 1),
                "tokens_generated": 0,
                "success": False,
            }

    def get_status(self) -> Dict:
        """后端状态"""
        return {
            "available_models": list(self._available.keys()),
            "loaded_models": list(self._loaded.keys()),
            "max_loaded": self.max_loaded,
            "device": str(self.device),
        }


# 全局单例
_backend = None
_backend_lock = threading.Lock()


def get_local_backend() -> LocalModelBackend:
    """获取本地推理后端单例"""
    global _backend
    with _backend_lock:
        if _backend is None:
            _backend = LocalModelBackend()
        return _backend


if __name__ == "__main__":
    # 测试
    backend = LocalModelBackend()
    print(f"可用模型: {backend.list_available()}")
    print(f"状态: {backend.get_status()}")
    
    for model_id in backend.list_available():
        print(f"\n--- 测试 {model_id} ---")
        result = backend.infer(model_id, "你好")
        print(f"  回复: {result['response'][:100]}")
        print(f"  延迟: {result['latency_ms']}ms")
        print(f"  成功: {result['success']}")