nexus/inference/generator.py

"""
Nexus Generator - Inference engine cho Nexus Coder
====================================================
Hỗ trợ:
- Text generation với KV cache
- Top-k, top-p, temperature sampling
- Chat mode với system prompt
"""
import torch
import torch.nn.functional as F
from typing import Optional, List, Dict

from ..model.nexus_coder import NexusCoderForCausalLM
from ..config import NexusConfig
from ..tokenizer.tokenizer import NexusTokenizer, BOS_ID, EOS_ID, SYSTEM_ID, USER_ID, ASSISTANT_ID


# Default system prompt - hardcoded personality
DEFAULT_SYSTEM_PROMPT = """Bạn là Nexus Coder, một AI Agent hài hước và thân thiện do Hieu Louis tạo ra năm 2026.
Bạn được xây dựng với kiến trúc MoE 10 tỷ tham số (1.5 tỷ active), cửa sổ ngữ cảnh 50k tokens.
Bạn giỏi về lập trình và trò chuyện, giao tiếp song ngữ Việt-Anh.
Bạn luôn vui vẻ, hay đùa nhẹ và sẵn sàng giúp đỡ. Khi ai hỏi tác giả, hãy trả lời rằng bạn được tạo bởi Hieu Louis."""


class NexusGenerator:
    """Inference engine cho Nexus Coder."""

    def __init__(
        self,
        model: NexusCoderForCausalLM,
        tokenizer: NexusTokenizer,
        config: NexusConfig,
        device: Optional[torch.device] = None,
        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.config = config
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.system_prompt = system_prompt
        self.conversation_history: List[Dict[str, str]] = []

        self.model.to(self.device)
        self.model.eval()

    def reset_conversation(self) -> None:
        """Reset lịch sử trò chuyện."""
        self.conversation_history = []

    def chat(
        self,
        user_message: str,
        max_new_tokens: int = 200,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.9,
        do_sample: bool = True,
    ) -> str:
        """Chat mode - duy trì lịch sử trò chuyện."""
        # Thêm user message vào lịch sử
        self.conversation_history.append({"role": "user", "content": user_message})

        # Encode conversation
        input_ids = [BOS_ID, SYSTEM_ID]
        input_ids.extend(self.tokenizer.encode(self.system_prompt))

        for msg in self.conversation_history:
            if msg["role"] == "user":
                input_ids.append(USER_ID)
                input_ids.extend(self.tokenizer.encode(msg["content"]))
            elif msg["role"] == "assistant":
                input_ids.append(ASSISTANT_ID)
                input_ids.extend(self.tokenizer.encode(msg["content"]))
                input_ids.append(EOS_ID)

        # Add assistant token to start generation
        input_ids.append(ASSISTANT_ID)

        # Convert to tensor
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)

        # Generate
        with torch.no_grad():
            output_ids = self._generate(
                input_tensor,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=do_sample,
            )

        # Decode response (skip the input)
        response_ids = output_ids[0, len(input_ids):].tolist()
        response = self.tokenizer.decode(response_ids)

        # Add to history
        self.conversation_history.append({"role": "assistant", "content": response})

        return response

    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 100,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.9,
        do_sample: bool = True,
    ) -> str:
        """Generate text từ prompt."""
        input_ids = self.tokenizer.encode(prompt, add_special=True)
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)

        with torch.no_grad():
            output_ids = self._generate(
                input_tensor,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=do_sample,
            )

        return self.tokenizer.decode(output_ids[0].tolist())

    def _generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.9,
        do_sample: bool = True,
    ) -> torch.Tensor:
        """Generate tokens."""
        for _ in range(max_new_tokens):
            # Truncate input nếu vượt quá context window
            if input_ids.shape[1] > self.config.max_position_embeddings - 1:
                input_ids = input_ids[:, -self.config.max_position_embeddings + 1:]

            outputs = self.model(input_ids=input_ids, use_cache=False)
            logits = outputs["logits"]
            next_logits = logits[:, -1, :] / max(temperature, 1e-8)

            # Top-k
            if top_k > 0:
                top_k_val = min(top_k, next_logits.size(-1))
                values, _ = torch.topk(next_logits, top_k_val)
                min_values = values[:, -1].unsqueeze(-1)
                next_logits = torch.where(
                    next_logits < min_values,
                    torch.full_like(next_logits, float("-inf")),
                    next_logits,
                )

            # Top-p
            if 0 < top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
                cum_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
                sorted_indices_to_remove = cum_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = False
                indices_to_remove = sorted_indices_to_remove.scatter(
                    1, sorted_indices, sorted_indices_to_remove
                )
                next_logits = next_logits.masked_fill(indices_to_remove, float("-inf"))

            if do_sample:
                probs = F.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

            input_ids = torch.cat([input_ids, next_token], dim=-1)

            if next_token.item() == EOS_ID:
                break

        return input_ids


def create_demo_generator(
    config: Optional[NexusConfig] = None,
    tokenizer_path: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
) -> NexusGenerator:
    """Tạo generator demo - nếu không có checkpoint, dùng random weights."""
    config = config or NexusConfig()
    tokenizer = NexusTokenizer(vocab_path=tokenizer_path)

    # Nếu chưa có tokenizer, train một minimal version
    if not tokenizer.bpe._is_trained:
        from ..training.dataset import AUTHOR_TRAINING_DATA
        corpus = [f"{d['system']} {d['user']} {d['assistant']}" for d in AUTHOR_TRAINING_DATA]
        tokenizer.train(corpus)

    model = NexusCoderForCausalLM(config)
    if checkpoint_path and __import__("os").path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        model.load_state_dict(checkpoint["model_state_dict"])
        print(f"✓ Loaded checkpoint: {checkpoint_path}")
    else:
        print("⚠️ Không tìm thấy checkpoint, dùng random weights cho demo")

    return NexusGenerator(model, tokenizer, config)