Update main.py

main.py

@@ -1,12 +1,15 @@
 """
 main.py - Entry point oficial para Hugging Face Spaces
-Inicializa AkiraAPI com persistência, logging limpo e health checks
+Inicializa AkiraAPI com persistência, logging limpo, health checks + HF LOGIN
 """
-
 import os
 import logging
+import shutil
 from modules.api import AkiraAPI
 import modules.config as config
+from loguru import logger
+from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 # ==========================
 # Configuração de Diretório Persistente (HF Spaces)
@@ -16,51 +19,92 @@ os.makedirs(DATA_DIR, exist_ok=True)
 config.DB_PATH = os.path.join(DATA_DIR, "akira.db")  # Banco persiste entre rebuilds!
 
 # ==========================
-# Logging (usando loguru do api.py, mas com fallback)
+# Logging (loguru com fallback)
 # ==========================
 try:
     from loguru import logger
 except ImportError:
-    # Fallback se loguru não estiver instalado
     logging.basicConfig(
         level=logging.INFO,
         format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
     )
     logger = logging.getLogger("AkiraMain")
 
+# ==========================
+# HF TOKEN LOGIN (SECRET NO HF SPACES)
+# ==========================
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN:
+    try:
+        login(token=HF_TOKEN)
+        logger.info("HF LOGIN OK COM TOKEN DO SECRET!")
+    except Exception as e:
+        logger.warning(f"HF LOGIN FALHOU: {e}")
+else:
+    logger.warning("HF_TOKEN NÃO ENCONTRADO → MODELOS PÚBLICOS OK, PRIVADOS NÃO!")
+
+# ==========================
+# VERIFICA / BAIXA HERMES 7B LOCAL (4-BIT)
+# ==========================
+MODEL_DIR = "/app/models/hermes-7b"
+FINETUNED_DIR = "/app/data/finetuned_hermes"
+
+def ensure_hermes_model():
+    """Baixa Hermes 7B 4-bit se não existir"""
+    if os.path.exists(f"{MODEL_DIR}/config.json"):
+        logger.info(f"Hermes 7B já existe em: {MODEL_DIR}")
+        return
+
+    logger.info("Hermes 7B NÃO ENCONTRADO → BAIXANDO 4-BIT...")
+    try:
+        quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+        tokenizer = AutoTokenizer.from_pretrained("NousResearch/Hermes-2-Pro-Mistral-7B")
+        model = AutoModelForCausalLM.from_pretrained(
+            "NousResearch/Hermes-2-Pro-Mistral-7B",
+            device_map="cpu",
+            quantization_config=quantization_config,
+            low_cpu_mem_usage=True,
+            torch_dtype="auto"
+        )
+        os.makedirs(MODEL_DIR, exist_ok=True)
+        tokenizer.save_pretrained(MODEL_DIR)
+        model.save_pretrained(MODEL_DIR)
+        logger.info("HERMES 7B 4-BIT BAIXADO E SALVO LOCALMENTE!")
+    except Exception as e:
+        logger.error(f"ERRO AO BAIXAR HERMES: {e}")
+        logger.error("Tente adicionar HF_TOKEN como Secret no HF Spaces!")
+
 # ==========================
 # Criação da Aplicação
 # ==========================
 def create_app():
-    """Inicializa a aplicação Akira IA com persistência e logging"""
-    logger.info("Iniciando Akira IA com Llama 3.1 8B (CPU 8-bit)...")
+    """Inicializa a aplicação Akira IA"""
+    logger.info("Iniciando Akira IA com Hermes 7B (CPU 4-bit)...")
     logger.info(f"Banco de dados: {config.DB_PATH}")
 
+    # Garante modelo local
+    ensure_hermes_model()
+
     try:
         akira = AkiraAPI(config)
         app = akira.app
-
         logger.info("AkiraAPI carregada com sucesso!")
         logger.info("Endpoints:")
-        logger.info("  POST /api/akira  → Chat com Akira")
-        logger.info("  GET  /health     → Health check")
-        logger.info("  GET  /           → Status")
-
+        logger.info(" POST /api/akira → Chat com Akira")
+        logger.info(" GET /health → Health check")
+        logger.info(" GET / → Status")
         return app
-
     except Exception as e:
         logger.error(f"Erro ao iniciar AkiraAPI: {e}")
         import traceback
         logger.error(traceback.format_exc())
         raise
 
-
 # ==========================
 # Instância da App
 # ==========================
 app = create_app()
 
-
 # ==========================
 # Rotas de Saúde
 # ==========================
@@ -69,13 +113,11 @@ def index():
     """Página inicial de status"""
     return "Akira IA online! Bué fixe, puto!", 200
 
-
 @app.route("/health")
 def health():
     """Health check para o container"""
     return "OK", 200
 
-
 # ==========================
 # Execução Local (modo dev)
 # ==========================