Spaces:

Shree-Lms-Developers2025
/

NovaAI

Running on Zero

NovaAI / src /model_engine.py

Shree-Lms-Developers2025

Deploy Nova AI (Flagship ChatGPT Astra Edition with NVIDIA ZeroGPU)

d58c800 verified 2 days ago

15.7 kB

"""
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

	"""
	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