src/model_engine.py

"""
Core Inference Engine for Nova AI - Flagship ChatGPT Astra Enterprise Suite.
Supports high-throughput streaming, multi-model switching, NVIDIA ZeroGPU & CUDA
hardware acceleration, BitsAndBytes 4-bit/8-bit quantization, PyTorch SDPA,
DeepSeek-R1 chain-of-thought, and resilient CPU fallback.
"""

import os
import sys
import threading
from pathlib import Path
from typing import Generator, List, Dict, Any, Optional, Tuple
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer
)

from src.config import config, AVAILABLE_MODELS, CACHE_DIR, EXPORTED_MODEL_DIR, GenerationConfig
from src.utils import get_system_hardware_info, LatencyTracker, console

# =====================================================================
# HUGGING FACE ZEROGPU (NVIDIA A100 DYNAMIC ACCELERATION) INTEGRATION
# =====================================================================
try:
    import spaces
    HAS_ZERO_GPU = True
except ImportError:
    HAS_ZERO_GPU = False
    class spaces:
        """Fallback decorator for local NVIDIA CUDA & non-ZeroGPU environments."""
        @staticmethod
        def GPU(func=None, duration=120):
            def decorator(f):
                return f
            if func is not None:
                return func
            return decorator


class NovaAstraEngine:
    """
    Flagship Enterprise inference engine tailored for Nova AI (Astra Edition).
    NVIDIA ZeroGPU & CUDA accelerated, supporting DeepSeek-R1, Qwen2.5, Llama-3.1,
    custom safetensors checkpoints, and deep chain-of-thought execution.
    """

    def __init__(self, model_id: Optional[str] = None):
        self.model_id = model_id or config.default_model_id
        self.device = self._resolve_device(config.device)
        self.model = None
        self.tokenizer = None
        self._is_loaded = False
        self._load_lock = threading.Lock()
        self.hardware_info = get_system_hardware_info()
        self._setup_cuda_optimizations()

    def _setup_cuda_optimizations(self):
        """Enable NVIDIA Tensor Core TF32 and CuDNN benchmark acceleration."""
        if torch.cuda.is_available():
            try:
                if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
                    torch.backends.cuda.matmul.allow_tf32 = True
                if hasattr(torch.backends, "cudnn"):
                    torch.backends.cudnn.benchmark = True
                    torch.backends.cudnn.allow_tf32 = True
            except Exception:
                pass

    def _resolve_device(self, preferred_device: str) -> str:
        """Resolve optimal compute device (NVIDIA CUDA / MPS / CPU)."""
        if preferred_device.lower() == "auto":
            if torch.cuda.is_available():
                return "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                return "mps"
            return "cpu"
        return preferred_device.lower()

    def load_model(self, model_id: Optional[str] = None, force_reload: bool = False):
        """Load or switch the active language model with NVIDIA CUDA / ZeroGPU optimization."""
        target_model_id = str(model_id or self.model_id)

        with self._load_lock:
            if self._is_loaded and target_model_id == str(self.model_id) and not force_reload:
                return

            console.print(f"[bold cyan]🌟 Loading Flagship Nova Astra Model:[/bold cyan] {target_model_id}")
            console.print(f"[dim]Hardware Target: {self.device.upper()} | ZeroGPU: {'Active' if HAS_ZERO_GPU else 'Native'} | Cache: {CACHE_DIR}[/dim]")

            CACHE_DIR.mkdir(parents=True, exist_ok=True)
            is_local_dir = Path(target_model_id).is_dir()

            # 1. Load Tokenizer
            tokenizer_kwargs = {"trust_remote_code": True}
            if not is_local_dir:
                tokenizer_kwargs["cache_dir"] = str(CACHE_DIR)

            try:
                self.tokenizer = AutoTokenizer.from_pretrained(target_model_id, **tokenizer_kwargs)
                if self.tokenizer.pad_token is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token
            except Exception as e:
                console.print(f"[yellow]Retrying tokenizer load with fast=False: {e}[/yellow]")
                tokenizer_kwargs["use_fast"] = False
                self.tokenizer = AutoTokenizer.from_pretrained(target_model_id, **tokenizer_kwargs)
                if self.tokenizer.pad_token is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token

            # 2. Build Model Configuration (NVIDIA CUDA / ZeroGPU / CPU)
            precision_mode = getattr(config.gen_config, "precision", "auto").lower()
            load_in_4bit = getattr(config.gen_config, "load_in_4bit", False) or (precision_mode == "4bit")
            load_in_8bit = (precision_mode == "8bit")

            model_kwargs: Dict[str, Any] = {
                "low_cpu_mem_usage": True,
                "trust_remote_code": True
            }
            if not is_local_dir:
                model_kwargs["cache_dir"] = str(CACHE_DIR)

            if self.device == "cuda":
                major_cc = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0
                total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3) if torch.cuda.is_available() else 0

                # Auto-detect if 4-bit is optimal (if VRAM is under 12GB for 7B+ models)
                if not load_in_4bit and not load_in_8bit and precision_mode == "auto" and total_vram_gb < 12.0:
                    load_in_4bit = True

                # Configure BitsAndBytes 4-bit / 8-bit quantization if requested or low VRAM
                if load_in_4bit or load_in_8bit:
                    try:
                        from transformers import BitsAndBytesConfig
                        if load_in_4bit:
                            compute_dtype = torch.bfloat16 if major_cc >= 8 else torch.float16
                            model_kwargs["quantization_config"] = BitsAndBytesConfig(
                                load_in_4bit=True,
                                bnb_4bit_compute_dtype=compute_dtype,
                                bnb_4bit_quant_type="nf4",
                                bnb_4bit_use_double_quant=True
                            )
                            console.print("[bold green]⚡ NVIDIA BitsAndBytes 4-bit NF4 Quantization Active (Max Speed, Low VRAM)[/bold green]")
                        elif load_in_8bit:
                            model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
                            console.print("[bold green]⚡ NVIDIA BitsAndBytes 8-bit Quantization Active[/bold green]")
                    except ImportError:
                        console.print("[dim]bitsandbytes not installed; continuing in native half-precision.[/dim]")

                # Optimal torch_dtype for NVIDIA Tensor Cores
                if "quantization_config" not in model_kwargs:
                    if precision_mode in ("bfloat16", "bf16") or (precision_mode == "auto" and major_cc >= 8):
                        model_kwargs["torch_dtype"] = torch.bfloat16
                    elif precision_mode in ("float16", "fp16") or precision_mode == "auto":
                        model_kwargs["torch_dtype"] = torch.float16
                    else:
                        model_kwargs["torch_dtype"] = torch.float32

                # Attention optimization via PyTorch SDPA (Scaled Dot-Product Attention)
                try:
                    model_kwargs["attn_implementation"] = "sdpa"
                except Exception:
                    pass

                model_kwargs["device_map"] = "auto"
            else:
                model_kwargs["torch_dtype"] = torch.float32

            # 3. Load Model with automatic fallback handling
            try:
                self.model = AutoModelForCausalLM.from_pretrained(target_model_id, **model_kwargs)
            except Exception as e:
                console.print(f"[yellow]Retrying model load without optional acceleration flags: {e}[/yellow]")
                model_kwargs.pop("attn_implementation", None)
                model_kwargs.pop("quantization_config", None)
                if self.device != "cuda":
                    model_kwargs.pop("device_map", None)
                self.model = AutoModelForCausalLM.from_pretrained(target_model_id, **model_kwargs)

            if "device_map" not in model_kwargs and self.device in ("cuda", "mps", "cpu"):
                try:
                    self.model.to(self.device)
                except Exception:
                    pass

            self.model.eval()
            self.model_id = target_model_id
            self._is_loaded = True
            console.print(f"[bold green]✓ Nova Astra Flagship '{target_model_id}' ready on {self.device.upper()}![/bold green]")

    def is_ready(self) -> bool:
        return self._is_loaded and self.model is not None and self.tokenizer is not None

    @spaces.GPU(duration=120)
    def stream_chat(
        self,
        messages: List[Dict[str, str]],
        gen_config: Optional[GenerationConfig] = None
    ) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
        """
        Stream generated tokens for multi-turn chat messages.
        Decorated with @spaces.GPU for seamless Hugging Face ZeroGPU NVIDIA A100 acceleration.
        Yields tuple: (accumulated_text, live_metrics_dict)
        """
        if not self.is_ready():
            self.load_model()

        cfg = gen_config or config.gen_config
        tracker = LatencyTracker()
        tracker.start()

        # Handle Deep Reasoning / Chain-of-Thought
        working_messages = list(messages)
        if getattr(cfg, "deep_reasoning", False):
            reasoning_instruction = (
                "\n\n[DEEP REASONING PROTOCOL ENABLED]:\n"
                "You are an advanced chain-of-thought reasoning intelligence. "
                "Analyze the problem systematically inside <think>...</think> (or <astra_thought>...</astra_thought>) tags. "
                "1. Break down the core objective and constraints.\n"
                "2. Evaluate hypotheses, mathematical logic, and edge cases.\n"
                "3. Verify your conclusions step-by-step.\n"
                "After closing the thought tag, provide the clean, authoritative, well-structured final answer."
            )
            if working_messages and working_messages[0]["role"] == "system":
                working_messages[0] = {
                    "role": "system",
                    "content": working_messages[0]["content"] + reasoning_instruction
                }
            else:
                working_messages.insert(0, {
                    "role": "system",
                    "content": "You are Nova AI (Astra Flagship Edition)." + reasoning_instruction
                })

        # Handle 1,000+ Domain Specialization Injection
        active_domain_name = None
        try:
            from src.domain_registry import domain_registry
            domain_id = getattr(cfg, "domain_id", None)
            target_domain = None
            if domain_id:
                target_domain = domain_registry.get_domain(domain_id)
            elif getattr(cfg, "auto_detect_domain", True):
                user_query = next((m.get("content", "") for m in reversed(working_messages) if m.get("role") == "user"), "")
                target_domain = domain_registry.auto_detect_domain(user_query)

            if target_domain:
                active_domain_name = f"{target_domain.name} ({target_domain.category})"
                domain_instruction = (
                    f"\n\n[ACTIVE DOMAIN SPECIALIZATION: {target_domain.name.upper()} ({target_domain.category})]:\n"
                    f"• Scope: {target_domain.description}\n"
                    f"• Authority Guidelines: {target_domain.guidelines}\n"
                    f"• Use domain-precise terminology, relevant formalisms, and industry best practices."
                )
                if working_messages and working_messages[0]["role"] == "system":
                    working_messages[0] = {
                        "role": "system",
                        "content": working_messages[0]["content"] + domain_instruction
                    }
                else:
                    working_messages.insert(0, {
                        "role": "system",
                        "content": "You are Nova AI (Astra Flagship Edition)." + domain_instruction
                    })
        except Exception:
            pass

        # Apply chat template
        try:
            formatted_prompt = self.tokenizer.apply_chat_template(
                working_messages,
                tokenize=False,
                add_generation_prompt=True
            )
        except Exception:
            formatted_prompt = ""
            for m in working_messages:
                role = m.get("role", "user").capitalize()
                formatted_prompt += f"{role}: {m.get('content', '')}\n"
            formatted_prompt += "Assistant: "

        # Direct tensors to appropriate device
        target_device = self.model.device if hasattr(self.model, "device") else self.device
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(target_device)
        input_token_count = inputs.input_ids.shape[-1]

        streamer = TextIteratorStreamer(
            self.tokenizer,
            timeout=60.0,
            skip_prompt=True,
            skip_special_tokens=True
        )

        max_tokens = int(getattr(cfg, "max_new_tokens", 1024) or 1024)
        generation_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=max_tokens,
            temperature=max(cfg.temperature, 0.01),
            top_p=cfg.top_p,
            repetition_penalty=cfg.repetition_penalty,
            do_sample=cfg.do_sample and cfg.temperature > 0.05,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self.tokenizer.eos_token_id
        )

        # Run token generation in background thread
        thread = threading.Thread(target=self.model.generate, kwargs=generation_kwargs)
        thread.start()

        accumulated_text = ""
        for new_token in streamer:
            tracker.record_token(1)
            accumulated_text += new_token
            current_metrics = tracker.stop()
            current_metrics["input_tokens"] = input_token_count
            if active_domain_name:
                current_metrics["active_domain"] = active_domain_name
            yield accumulated_text, current_metrics

        thread.join()
        final_metrics = tracker.stop()
        final_metrics["input_tokens"] = input_token_count
        if active_domain_name:
            final_metrics["active_domain"] = active_domain_name
        yield accumulated_text, final_metrics

    @spaces.GPU(duration=120)
    def generate_response(
        self,
        messages: List[Dict[str, str]],
        gen_config: Optional[GenerationConfig] = None
    ) -> Tuple[str, Dict[str, Any]]:
        """Non-streaming generation helper."""
        full_text = ""
        metrics = {}
        for partial, cur_metrics in self.stream_chat(messages, gen_config):
            full_text = partial
            metrics = cur_metrics
        return full_text, metrics


# Global shared engine instance
engine = NovaAstraEngine()

# Enterprise & Backwards compatibility aliases
FlagshipChatEngine = NovaAstraEngine
LightweightChatEngine = NovaAstraEngine