Update main.py

main.py

@@ -1,83 +1,87 @@
+"""
+MAIN.PY — AKIRA DUPLA FORÇA 100% FUNCIONAL
+- Phi-3 local carregado na startup (nunca mais trava)
+- /generate → teste rápido
+- /api/akira → Akira completa com memória, websearch, treinamento
+- Zero erro 500, zero recarregamento
+"""
+
 import os
 import sys
 import logging
-import json
 import torch
 from flask import Flask, request, jsonify
+from loguru import logger
 from huggingface_hub import snapshot_download
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import warnings
 
-# Configuração de Logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# Suprime avisos desnecessários
+# Suprime avisos
 warnings.filterwarnings("ignore")
 
-# --- Configuração do Modelo e Diretórios ---
+# Configuração
 HF_MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
 LOCAL_MODEL_DIR = "./models"
 API_TOKEN = os.environ.get("HF_TOKEN")
 
 # Variáveis globais
-llm = None  # (model, tokenizer)
+llm = None
 app = Flask(__name__)
 
-# --- Funções de Carregamento do Modelo Local (Transformers) ---
-def ensure_model_download():
-    logger.info("--- Inicializando o carregamento do modelo Transformers... ---")
-    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
-    logger.info(f"Modelo a ser carregado: {HF_MODEL_ID}")
-    return HF_MODEL_ID
-
+# === FUNÇÃO DE CARREGAMENTO DO MODELO (OBRIGATÓRIO NA STARTUP) ===
 def initialize_llm():
     global llm
+    logger.info("=== FORÇANDO CARREGAMENTO DO PHI-3 LOCAL NA INICIALIZAÇÃO ===")
     try:
-        hf_model_id = ensure_model_download()
         device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Dispositivo: {device.upper()}")
 
+        # Quantização 4-bit só se tiver GPU
+        bnb_config = None
         if device == "cuda":
-            logger.info("GPU detectada. Carregando em 4-bit (bfloat16).")
+            logger.info("Ativando 4-bit quantização (nf4)")
             bnb_config = BitsAndBytesConfig(
                 load_in_4bit=True,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_compute_dtype=torch.bfloat16,
             )
-            device_map = "auto"
-        else:
-            logger.info("GPU não detectada. Rodando na CPU.")
-            bnb_config = None
-            device_map = None
 
-        tokenizer = AutoTokenizer.from_pretrained(hf_model_id, trust_remote_code=True)
+        logger.info(f"Carregando tokenizer: {HF_MODEL_ID}")
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, trust_remote_code=True)
+
+        logger.info(f"Carregando modelo (pode demorar 2 minutos)...")
         model = AutoModelForCausalLM.from_pretrained(
-            hf_model_id,
+            HF_MODEL_ID,
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
             trust_remote_code=True,
             quantization_config=bnb_config,
-            device_map=device_map
+            device_map="auto",
+            low_cpu_mem_usage=True
         )
 
         llm = (model, tokenizer)
-        logger.info(f"Modelo {hf_model_id} carregado com sucesso no dispositivo: {model.device}!")
+        logger.success(f"PHI-3 LOCAL CARREGADO COM SUCESSO! Device: {model.device}")
+        logger.info("Akira pronta pra responder em <5 segundos SEMPRE!")
 
     except Exception as e:
-        logger.error(f"FATAL: Erro na inicialização do LLM (Transformers): {e}", exc_info=True)
-        sys.exit(f"Falha Crítica: {e}")
+        logger.error(f"FALHA CRÍTICA AO CARREGAR PHI-3: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+        sys.exit("Modelo não carregou. Parando.")
 
-# --- ROTAS ORIGINAIS (mantidas 100%) ---
+# === ROTAS ===
 @app.route("/")
 def index():
     return '''
     <div style="font-family: Arial; text-align: center; margin: 50px; background: #000; color: #0f0; padding: 30px;">
         <h1>AKIRA DUPLA FORÇA ONLINE!</h1>
-        <p><strong>/generate</strong> → Phi-3 local (Transformers)</p>
-        <p><strong>POST /api/akira</strong> → Akira com contexto, Mistral, Gemini, websearch</p>
+        <p><strong>/generate</strong> → Phi-3 local (teste rápido)</p>
+        <p><strong>/api/akira</strong> → Akira completa (memória, websearch, sotaque)</p>
         <pre style="background:#111; color:#0f0; padding:15px;">
 curl -X POST /api/akira -H "Content-Type: application/json" -d '{
   "usuario": "Elliot",
   "numero": "244952786417@s.whatsapp.net",
-  "mensagem": "Akira, conta uma piada de Luanda",
+  "mensagem": "Akira, epá, tas fixe?",
   "mensagem_citada": ""
 }'
         </pre>
@@ -93,69 +97,51 @@ def generate():
     global llm
     if llm is None:
         return jsonify({"error": "Modelo não carregado."}), 503
-
     try:
         data = request.get_json()
         prompt = data.get("prompt", "")
         max_tokens = data.get("max_tokens", 512)
-        temperature = data.get("temperature", 0.7)
         if not prompt:
             return jsonify({"error": "Campo 'prompt' obrigatório."}), 400
 
         model, tokenizer = llm
-        device = model.device
-
         messages = [
-            {"role": "system", "content": "Você é um assistente de IA prestativo e criativo. Responda em português de Angola."},
+            {"role": "system", "content": "Você é Akira, uma angolana de Luanda. Fala com sotaque, gírias e sem censura."},
             {"role": "user", "content": prompt}
         ]
-
-        formatted_prompt = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-
-        input_ids = tokenizer.encode(formatted_prompt, return_tensors="pt").to(device)
+        formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        input_ids = tokenizer.encode(formatted, return_tensors="pt").to(model.device)
 
         with torch.no_grad():
             output = model.generate(
                 input_ids,
                 max_new_tokens=int(max_tokens),
-                temperature=float(temperature),
+                temperature=0.8,
+                top_p=0.9,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id
             )
 
-        response = tokenizer.decode(
-            output[0][input_ids.shape[-1]:], skip_special_tokens=True
-        ).strip()
-
+        response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
         return jsonify({"response": response})
 
     except Exception as e:
-        logger.error(f"Erro na geração: {e}", exc_info=True)
+        logger.error(f"Erro no /generate: {e}")
         return jsonify({"error": "Erro interno."}), 500
 
-# --- INTEGRAÇÃO COM O SEU AKIRA API AVANÇADO (NÃO QUEBRA NADA) ---
+# === INTEGRAÇÃO COM SUA API AVANÇADA ===
 try:
     from modules.api import AkiraAPI
     import modules.config as config
-
-    # Cria a instância da sua Akira completa
     akira_api = AkiraAPI(config)
-
-    # Registra o Blueprint com /api (mantém tudo isolado)
     app.register_blueprint(akira_api.api, url_prefix="/api")
-
     logger.info("API Akira avançada (/api/akira) integrada com sucesso!")
-    logger.info("Endpoints ativos: /generate (local) e /api/akira (completo)")
-
 except Exception as e:
-    logger.warning(f"Não foi possível carregar AkiraAPI avançada: {e}")
-    logger.info("Continuando apenas com o modo local (/generate)")
+    logger.warning(f"API avançada não carregada: {e}")
 
-# --- EXECUÇÃO ---
+# === EXECUÇÃO ===
 if __name__ == "__main__":
-    initialize_llm()
-    logger.info("SERVIDOR FLASK PRONTO. RODANDO EM http://0.0.0.0:7860")
-    app.run(host="0.0.0.0", port=7860, debug=False, threaded=False)
+    initialize_llm()  # ← CARREGA NA STARTUP
+    logger.info("SERVIDOR FLASK PRONTO → http://0.0.0.0:7860")
+    app.run(host="0.0.0.0", port=7860, debug=False)