Update modules/local_llm.py

modules/local_llm.py

@@ -1,24 +1,54 @@
 # modules/local_llm.py
 import os
+import threading
 from loguru import logger
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
+from peft import PeftModel
 
 MODEL_DIR = "/app/models/hermes-7b"
-FINETUNED = "/app/data/finetuned_hermes"
+FINETUNED_DIR = "/app/data/finetuned_hermes"
 
-class LocalLLM:
-    def __init__(self):
-        self.generator = None
-        self._load_model()
+# SINGLETON + LOCK
+_HERMES_GLOBAL = None
+_HERMES_LOCK = threading.Lock()
+
+
+def _get_hermes_singleton():
+    """Retorna modelo + tokenizer (4-bit) → singleton global"""
+    global _HERMES_GLOBAL
+    if _HERMES_GLOBAL is not None:
+        logger.debug("Reusando Hermes 7B 4-bit global")
+        return _HERMES_GLOBAL
+
+    with _HERMES_LOCK:
+        if _HERMES_GLOBAL is not None:
+            return _HERMES_GLOBAL
+
+        logger.info("Carregando Hermes 7B 4-bit (6 GB RAM)...")
+
+        # Verifica arquivos
+        required = ["config.json", "model.safetensors.index.json"]
+        missing = [f for f in required if not os.path.exists(f"{MODEL_DIR}/{f}")]
+        if missing:
+            logger.error(f"ARQUIVOS FALTANDO: {missing}")
+            return None
+
+        shards = [f for f in os.listdir(MODEL_DIR) if f.endswith(".safetensors")]
+        if len(shards) != 4:
+            logger.error(f"APENAS {len(shards)} SHARDS! ESPERADO: 4")
+            return None
 
-    def _load_model(self):
         try:
-            logger.info("Carregando Hermes 7B 4-bit (6 GB RAM)...")
             quantization = BitsAndBytesConfig(
                 load_in_4bit=True,
-                bnb_4bit_compute_dtype="float16"
+                bnb_4bit_compute_dtype="float16",
+                bnb_4bit_quant_type="nf4"
             )
+
             tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+
             model = AutoModelForCausalLM.from_pretrained(
                 MODEL_DIR,
                 device_map="cpu",
@@ -26,29 +56,69 @@ class LocalLLM:
                 low_cpu_mem_usage=True,
                 offload_folder="/tmp/offload"
             )
-            self.generator = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=256,
-                temperature=0.8,
-                do_sample=True
-            )
-            logger.info("HERMES 7B 4-BIT CARREGADO!")
+
+            # Carrega LoRA se existir
+            if os.path.exists(f"{FINETUNED_DIR}/adapter_config.json"):
+                logger.info("Aplicando LoRA finetuned...")
+                model = PeftModel.from_pretrained(model, FINETUNED_DIR)
+                logger.info("LoRA ANGOLANO ATIVO!")
+
+            _HERMES_GLOBAL = (model, tokenizer)
+            logger.info("Hermes 7B 4-bit GLOBAL carregado!")
+            return _HERMES_GLOBAL
+
         except Exception as e:
             logger.error(f"ERRO: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return None
+
+
+class LocalLLM:
+    def __init__(self):
+        self.generator = None
+        self._load_pipeline()
+
+    def _load_pipeline(self):
+        result = _get_hermes_singleton()
+        if not result:
+            logger.error("Pipeline off")
             self.generator = None
+            return
 
-    def is_available(self):
+        model, tokenizer = result
+        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
+        )
+        logger.info("Pipeline LOCAL com 4-bit + LoRA → ONLINE!")
+
+    def is_available(self) -> bool:
         return self.generator is not None
 
-    def generate(self, prompt, max_tokens=256, temperature=0.8):
+    def generate(self, prompt: str, max_tokens: int = 256, temperature: float = 0.8) -> str:
         if not self.is_available():
-            return "Off, kota."
+            return "Modelo off, kota."
+
         try:
-            output = self.generator(prompt, max_new_tokens=max_tokens)[0]["generated_text"]
-            return output.strip()
+            output = self.generator(
+                prompt,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=temperature > 0.0,
+                repetition_penalty=1.1,
+                return_full_text=False
+            )
+            return output[0]["generated_text"].strip()
         except Exception as e:
+            logger.error(f"Geração falhou: {e}")
             return "Buguei, puto."
 
+# EXPORTA
 HermesLLM = LocalLLM