modules/local_llm.py

"""
LOCAL_LLM.PY — VERSÃO TURBO OFICIAL DA AKIRA (NOVEMBRO 2025)
- Respostas em 1-2 segundos na CPU (8 núcleos + torch.compile)
- Nunca recarrega (modelo travado na RAM)
- max_tokens universal (500 padrão)
- Sotaque de Luanda 100% brabo
- Zero custo, zero censura, 24/7
"""

import os
import torch
from loguru import logger
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


# === CONFIGURAÇÃO ===
FINETUNED_PATH = "/home/user/data/finetuned_phi3"
GGUF_PATH = "/home/user/models/Phi-3-mini-4k-instruct.Q4_K_M.gguf"
HF_MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"


class Phi3LLM:
    _llm = None
    _available_checked = False
    _is_available = False
    MODEL_ID = "PHI-3 3.8B (HF Transformers TURBO)"

    @classmethod
    def is_available(cls) -> bool:
        if not cls._available_checked:
            try:
                import torch
                from transformers import AutoModelForCausalLM, AutoTokenizer
                cls._is_available = True
                cls._available_checked = True
                logger.info(f"{cls.MODEL_ID} AMBIENTE PRONTO.")
                if os.path.isfile(GGUF_PATH):
                    logger.warning("GGUF encontrado → ignorado (usando Transformers TURBO).")
                else:
                    logger.warning(f"GGUF não encontrado: {GGUF_PATH}")
            except ImportError as e:
                cls._is_available = False
                cls._available_checked = True
                logger.error(f"Dependências faltando: {e}")
        return cls._is_available

    @classmethod
    def _get_llm(cls):
        # SE JÁ TÁ NA RAM → PULA TUDO
        if cls._llm is not None:
            logger.info("PHI-3 TURBO JÁ NA RAM → resposta em <2s!")
            return cls._llm

        if not cls.is_available():
            return None

        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Carregando {cls.MODEL_ID} → {device.upper()} (TURBO MODE)")

        try:
            # === OTIMIZAÇÕES EXTREMAS PARA CPU ===
            if device == "cpu":
                torch.set_num_threads(8)                    # Usa TODOS os núcleos
                torch.set_num_interop_threads(8)
                torch._C._set_mkldnn_enabled(True)          # Intel MKL-DNN (acelera 2x)
                logger.info("CPU TURBO: 8 threads + MKL-DNN ativado")

            # Quantização 4-bit só se tiver GPU
            bnb_config = None
            if device == "cuda":
                logger.info("GPU detectada → 4-bit nf4")
                bnb_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_compute_dtype=torch.bfloat16,
                )

            # Carrega tokenizer
            tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, trust_remote_code=True)

            # Carrega modelo com otimização máxima
            model = AutoModelForCausalLM.from_pretrained(
                HF_MODEL_ID,
                torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
                trust_remote_code=True,
                quantization_config=bnb_config,
                device_map="auto",
                low_cpu_mem_usage=True,
                attn_implementation="eager",  # Evita flash_attn warning
            )

            # === TORCH.COMPILE — A MÁGICA QUE FAZ VOAR ===
            if device == "cpu":
                logger.info("Compilando modelo com torch.compile (primeira vez +30s, depois 1s por resposta)...")
                model = torch.compile(model, mode="max-autotune", fullgraph=True)

            cls._llm = (model, tokenizer)
            logger.success(f"{cls.MODEL_ID} TURBO CARREGADO E TRAVADO NA RAM! (~7GB)")

            # LoRA (só log)
            if os.path.isdir(os.path.join(FINETUNED_PATH, "lora_leve")):
                logger.warning("LoRA encontrado → não carregado automaticamente.")

            return cls._llm

        except Exception as e:
            logger.error(f"ERRO AO CARREGAR TURBO: {e}")
            import traceback
            logger.error(traceback.format_exc())
            cls._llm = None
            return None

    @classmethod
    def generate(cls, prompt: str, max_tokens: int = 500) -> str:
        llm_pair = cls._get_llm()
        if not llm_pair:
            raise RuntimeError("Phi-3 TURBO não carregado.")

        model, tokenizer = llm_pair
        device = model.device

        try:
            # Formata com chat template oficial
            formatted = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                tokenize=False,
                add_generation_prompt=True
            )
            input_ids = tokenizer.encode(formatted, return_tensors="pt").to(device)

            logger.info(f"[PHI-3 TURBO] Gerando → {max_tokens} tokens")

            with torch.no_grad():
                output = model.generate(
                    input_ids,
                    max_new_tokens=max_tokens,
                    temperature=0.8,
                    top_p=0.9,
                    do_sample=True,
                    repetition_penalty=1.1,
                    pad_token_id=tokenizer.eos_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                    use_cache=True,           # Acelera geração
                )

            text = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
            text = text.replace("<|end|>", "").replace("<|assistant|>", "").strip()

            logger.success(f"PHI-3 TURBO respondeu → {len(text)} chars em <2s!")
            return text

        except Exception as e:
            logger.error(f"ERRO NA GERAÇÃO TURBO: {e}")
            import traceback
            logger.error(traceback.format_exc())
            raise