Update modules/local_llm.py

modules/local_llm.py

@@ -1,71 +1,107 @@
-# modules/local_llm.py (ATUALIZADO + EXPORT CERTO)
+# modules/local_llm.py
 import os
 from loguru import logger
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
+# === CONFIGURAÇÕES ===
 MODEL_DIR = "/app/models/hermes-7b"
+FINETUNED_PATH = "/app/data/finetuned_hermes"
+
+# === SINGLETON GLOBAL (COMPARTILHADO COM treinamento.py) ===
+_HERMES_GLOBAL = None
+
+
+def _get_hermes_singleton():
+    """Carrega ou retorna Hermes com LoRA (singleton global)"""
+    global _HERMES_GLOBAL
+    if _HERMES_GLOBAL is not None:
+        logger.debug("Reusando Hermes 7B global (local_llm)")
+        return _HERMES_GLOBAL
+
+    logger.info("Carregando Hermes 7B UMA VEZ (local + finetune)...")
+
+    try:
+        if not os.path.exists(f"{MODEL_DIR}/config.json"):
+            logger.error("config.json NÃO ENCONTRADO!")
+            return None
+
+        # Verifica shards
+        shards = [f for f in os.listdir(MODEL_DIR) if f.endswith(".safetensors")]
+        if len(shards) != 4:
+            logger.error(f"APENAS {len(shards)} SHARDS .safetensors! FALTANDO!")
+            return None
+
+        # Carrega tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        # Carrega base model (CPU, low mem)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_DIR,
+            device_map="cpu",
+            torch_dtype="float16",
+            low_cpu_mem_usage=True,
+            offload_folder="/tmp/offload",
+            offload_state_dict=True
+        )
+
+        # Carrega LoRA se existir
+        if os.path.exists(f"{FINETUNED_PATH}/adapter_config.json"):
+            from peft import PeftModel
+            logger.info("Carregando LoRA finetuned...")
+            model = PeftModel.from_pretrained(model, FINETUNED_PATH)
+            logger.info("LoRA ANGOLANO CARREGADO!")
+        else:
+            logger.info("Nenhum LoRA encontrado. Usando base.")
+
+        _HERMES_GLOBAL = (model, tokenizer)
+        logger.info("Hermes 7B GLOBAL carregado com sucesso!")
+        return _HERMES_GLOBAL
+
+    except Exception as e:
+        logger.error(f"Erro ao carregar Hermes global: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+        return None
+
 
 class LocalLLM:
     def __init__(self):
         self.generator = None
-        self._load_model()
+        self._load_pipeline()
 
-    def _load_model(self):
-        try:
-            logger.info(f"Carregando Hermes 7B 8-bit (CPU) de: {MODEL_DIR}")
-
-            if not os.path.exists(f"{MODEL_DIR}/config.json"):
-                logger.error(f"MODELO NÃO ENCONTRADO EM: {MODEL_DIR}")
-                self.generator = None
-                return
-
-            if not any(f.endswith((".safetensors", ".bin")) for f in os.listdir(MODEL_DIR)):
-                logger.error("NENHUM SHARD ENCONTRADO! RECONSTRUIR COM DOCKERFILE!")
-                self.generator = None
-                return
-
-            offload_dir = "/tmp/offload"
-            os.makedirs(offload_dir, exist_ok=True)
-
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)
-
-            model = AutoModelForCausalLM.from_pretrained(
-                MODEL_DIR,
-                device_map="cpu",
-                torch_dtype="float16",
-                low_cpu_mem_usage=True,
-                offload_folder=offload_dir,
-                offload_state_dict=True
-            )
+    def _load_pipeline(self):
+        result = _get_hermes_singleton()
+        if not result:
+            logger.error("Hermes não carregado. Pipeline indisponível.")
+            self.generator = None
+            return
 
-            self.generator = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=256,
-                temperature=0.8,
-                do_sample=True,
-                repetition_penalty=1.1,
-                return_full_text=False
-            )
+        model, tokenizer = result
 
-            logger.info("HERMES 7B 8-BIT (CPU) CARREGADO COM SUCESSO!")
-        except Exception as e:
-            logger.error(f"ERRO AO CARREGAR MODELO: {e}")
-            import traceback
-            logger.error(traceback.format_exc())
-            self.generator = None
+        self.generator = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=256,
+            temperature=0.8,
+            do_sample=True,
+            repetition_penalty=1.1,
+            return_full_text=False,
+            device_map="cpu"
+        )
+        logger.info("Pipeline LOCAL conectado ao Hermes com LoRA!")
 
     def is_available(self) -> bool:
         return self.generator is not None
 
     def generate(self, prompt: str, max_tokens: int = 256, temperature: float = 0.8) -> str:
         if not self.is_available():
-            return "Desculpa, kota... o modelo não carregou."
+            return "Desculpa, kota... o modelo tá off."
 
         try:
-            logger.info(f"[HERMES] Gerando com max_tokens={max_tokens}, temp={temperature}")
-
+            logger.info(f"[LOCAL] Gerando: max_tokens={max_tokens}, temp={temperature}")
             output = self.generator(
                 prompt,
                 max_new_tokens=max_tokens,
@@ -75,11 +111,11 @@ class LocalLLM:
                 return_full_text=False
             )
             text = output[0]["generated_text"].strip()
-            logger.info(f"[HERMES] Resposta: {text[:60]}...")
+            logger.info(f"[LOCAL] Resposta: {text[:60]}...")
             return text
         except Exception as e:
-            logger.error(f"ERRO NA GERAÇÃO LOCAL: {e}")
+            logger.error(f"Erro na geração local: {e}")
             return "Bué, deu pau no Hermes local..."
 
-# EXPORTA COM NOME QUE O api.py ESPERA
+# EXPORTA PARA api.py
 HermesLLM = LocalLLM