server.py

#!/usr/bin/env python3
"""
Wrapper для обратной совместимости.

Запускает EmbryonicServer (пренатальная фаза).
После рождения модели переключается на InferenceServer.

Usage:
    python3 server.py --config server_config.yaml --port 8080
"""

import sys
from pathlib import Path
project_root = str(Path(__file__).parent.resolve())
if project_root not in sys.path:
    sys.path.insert(0, project_root)

import argparse
import os


class SimpleBirthController:
    """
    Простой контроллер для загрузки модели с адаптером.
    
    Используется в embryonic сервере для вычисления real surprise.
    """
    
    def __init__(self, model_path: str, adapter_path: str, config_path: str):
        from utils.config import SomnolentConfig
        from adapters.qwen35_adapter import Qwen35Adapter, Qwen35Config
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        
        print("  Загрузка токенизатора...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True,
        )
        
        print("  Загрузка модели...")
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map="cuda:0" if torch.cuda.is_available() else "cpu",
        )
        
        print("  Создание адаптера...")
        somnolent_config = SomnolentConfig.from_yaml(config_path)
        lora_config = somnolent_config.lora_config
        
        adapter_config = Qwen35Config(
            r=lora_config.get("r", 32),
            lora_alpha=lora_config.get("lora_alpha", 64),
            lora_dropout=lora_config.get("lora_dropout", 0.05),
            target_modules=lora_config.get("target_modules", [
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"
            ]),
            init_method="svd",
            model_path=model_path,
            transformer_layer_indices=somnolent_config.transformer_layer_indices,
        )
        
        self.adapter = Qwen35Adapter(
            config=adapter_config,
            model=self.model,
            verbose=False,
        )
        
        print(f"  Загрузка весов адаптера из {adapter_path}...")
        self.adapter.load_adapter_state(torch.load(adapter_path, map_location=self.model.device))
        
        # Разморозка адаптера для обучения
        self.adapter.unfreeze_adapter()
        
        # Создаём episodic buffer
        from memory.episodic_buffer import EpisodicBuffer, EpisodicBufferConfig
        buffer_config = EpisodicBufferConfig(max_capacity=10000)
        self.episodic_buffer = EpisodicBuffer(config=buffer_config)
        
        print(f"  Модель готова: {self.model.device}")
    
    async def awaken(self):
        """Заглушка для совместимости."""
        return self.model


def main():
    parser = argparse.ArgumentParser(description="SomnolentLLM Server")
    parser.add_argument("--config", type=str, default="server_config.yaml", help="Путь к конфигу")
    parser.add_argument("--port", type=int, default=None, help="Port для сервера")
    parser.add_argument("--host", type=str, default=None, help="Host для сервера")
    parser.add_argument("--mode", type=str, default=None, choices=["embryonic", "inference"],
                        help="Режим сервера (embryonic/inference)")
    args = parser.parse_args()

    # Загрузка конфигурации
    config_path = args.config
    if not os.path.exists(config_path):
        print(f"❌ Конфигурация не найдена: {config_path}")
        sys.exit(1)

    import yaml
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    # Определение режима
    embryonic_config = config.get("embryonic_mode", {})
    embryonic_enabled = embryonic_config.get("enabled", False)

    # Переопределение режима из аргументов
    mode = args.mode
    if mode is None:
        # Авто-определение: если embryonic_enabled → embryonic, иначе inference
        mode = "embryonic" if embryonic_enabled else "inference"

    # Переопределение порта
    port = args.port
    if port is None:
        if mode == "embryonic":
            port = embryonic_config.get("port", 8080)
        else:
            port = config.get("port", 8000)

    host = args.host or config.get("host", "0.0.0.0")

    print("=" * 60)
    print("SomnolentLLM Server")
    print("=" * 60)
    print(f"Mode: {mode}")
    print(f"Host: {host}")
    print(f"Port: {port}")
    print("=" * 60)

    # Запуск соответствующего сервера
    if mode == "embryonic":
        print("\n🤰 Запуск EmbryonicServer...")
        from servers.embryonic_server import EmbryonicServer, EmbryonicServerConfig
        from utils.config import SomnolentConfig

        server_config = EmbryonicServerConfig(
            host=host,
            port=port,
            auto_birth=embryonic_config.get("auto_birth", True),
            min_gestation_hours=embryonic_config.get("min_gestation_hours", 1.0),
            max_gestation_hours=embryonic_config.get("max_gestation_hours", 12.0),
            max_memories=embryonic_config.get("max_memories", 10000),
            entropy_threshold=embryonic_config.get("entropy_threshold", 0.1),
        )

        # Загрузка модели для реального вычисления surprise
        adapter_path = config.get("adapter_path", "./output/born_adapter_v5.pt")
        model_config_path = config.get("model_config", "config.qwen35.yaml")
        
        print(f"\n🧬 Загрузка модели...")
        print(f"  Config: {model_config_path}")
        print(f"  Adapter: {adapter_path}")
        
        try:
            # Получаем путь к модели из конфига
            somnolent_config = SomnolentConfig.from_yaml(model_config_path)
            model_path = somnolent_config.model_path
            
            birth_controller = SimpleBirthController(
                model_path=model_path,
                adapter_path=adapter_path,
                config_path=model_config_path,
            )
            print("✅ Модель загружена")
        except Exception as e:
            print(f"⚠️  Не удалось загрузить модель: {e}")
            print("  Запуск без модели (дефолтный surprise)")
            birth_controller = None

        server = EmbryonicServer(
            config=server_config,
            birth_controller=birth_controller,
            verbose=True,
        )
        server.run()

    else:
        print("\n🧠 Запуск InferenceServer...")
        from servers.inference_server import InferenceServer, InferenceServerConfig

        server_config = InferenceServerConfig(
            host=host,
            port=port,
            model_config=config.get("model_config", "config.qwen35.yaml"),
            adapter_path=config.get("adapter_path"),
            fatigue_enabled=config.get("fatigue", {}).get("enabled", True),
            fatigue_check_interval=config.get("fatigue", {}).get("check_interval_sec", 60),
            fatigue_threshold=config.get("fatigue", {}).get("fatigue_threshold", 0.7),
            auto_sleep=config.get("fatigue", {}).get("auto_sleep", True),
            nrem_steps=config.get("sleep", {}).get("nrem_steps", 30),
            rem_steps=config.get("sleep", {}).get("rem_steps", 10),
            min_sleep_cycles=config.get("sleep", {}).get("sleep_cycles", 2),
        )

        server = InferenceServer(config=server_config, verbose=True)
        server.run()


if __name__ == "__main__":
    main()