Update modules/local_llm.py

modules/local_llm.py

@@ -1,107 +1,72 @@
 # modules/local_llm.py
 import os
-import threading
 from loguru import logger
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from peft import PeftModel
-
-MODEL_DIR = "/app/models/hermes-7b"
-FINETUNED_DIR = "/app/data/finetuned_hermes"
-
-_HERMES_GLOBAL = None
-_HERMES_LOCK = threading.Lock()
-
+from llama_cpp import Llama
+import threading
 
-def _get_hermes_singleton():
-    global _HERMES_GLOBAL
-    if _HERMES_GLOBAL is not None:
-        logger.debug("Reusando Hermes 7B FP16 global")
-        return _HERMES_GLOBAL
+MODEL_PATH = "/app/models/openhermes-2.5-mistral-7b.Q4_K_M.gguf"
+FINETUNED_PATH = "/app/data/finetuned_hermes"
 
-    with _HERMES_LOCK:
-        if _HERMES_GLOBAL is not None:
-            return _HERMES_GLOBAL
+_llm_global = None
+_lock = threading.Lock()
 
-        logger.info("Carregando Hermes 7B FP16 + offload (8 GB RAM)...")
+def _get_llm():
+    global _llm_global
+    if _llm_global is not None:
+        return _llm_global
 
-        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"FALTANDO: {missing}")
-            return None
+    with _lock:
+        if _llm_global is not None:
+            return _llm_global
 
-        shards = [f for f in os.listdir(MODEL_DIR) if f.endswith(".safetensors")]
-        if len(shards) != 4:
-            logger.error(f"SHARDS: {len(shards)}/4")
+        logger.info("Carregando OpenHermes-2.5-Mistral-7B Q4_K_M (4.8 GB RAM)...")
+        if not os.path.exists(MODEL_PATH):
+            logger.error("GGUF não encontrado! Baixa com snapshot_download.")
             return None
 
         try:
-            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,
-                torch_dtype="auto",           # fp16 se tiver GPU, senão bfloat16
-                device_map="cpu",
-                low_cpu_mem_usage=True,
-                offload_folder="/tmp/offload",
-                offload_state_dict=True
+            llm = Llama(
+                model_path=MODEL_PATH,
+                n_ctx=4096,
+                n_threads=4,
+                n_gpu_layers=0,  # CPU only
+                n_batch=512,
+                verbose=False
             )
 
-            if os.path.exists(f"{FINETUNED_DIR}/adapter_config.json"):
-                logger.info("Carregando LoRA angolano...")
-                model = PeftModel.from_pretrained(model, FINETUNED_DIR)
-                logger.info("LoRA ATIVO!")
-
-            _HERMES_GLOBAL = (model, tokenizer)
-            logger.info("Hermes 7B FP16 GLOBAL carregado com sucesso!")
-            return _HERMES_GLOBAL
+            if os.path.exists(f"{FINETUNED_PATH}/adapter_config.json"):
+                logger.info("LoRA angolano DETECTADO → aplicando...")
+                llm.load_lora(FINETUNED_PATH)
 
+            _llm_global = llm
+            logger.info("OpenHermes 2.5 Q4_K_M + LoRA → ONLINE EM 4.8 GB!")
+            return llm
         except Exception as e:
-            logger.error(f"ERRO CRÍTICO: {e}")
-            import traceback
-            logger.error(traceback.format_exc())
+            logger.error(f"ERRO GGUF: {e}")
             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("Hermes off → usando API")
-            self.generator = None
-            return
-
-        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 FP16 + LoRA → ONLINE!")
+        self.llm = _get_llm()
 
     def is_available(self) -> bool:
-        return self.generator is not None
+        return self.llm is not None
 
-    def generate(self, prompt: str, max_tokens: int = 256, temperature: float = 0.8) -> str:
+    def generate(self, prompt: str, max_tokens: int = 256) -> str:
         if not self.is_available():
-            return "Modelo local off, kota."
-
+            return "Modelo local off, usando API."
         try:
-            out = self.generator(prompt, max_new_tokens=max_tokens, temperature=temperature)[0]["generated_text"]
-            return out.strip()
+            output = self.llm(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=0.8,
+                top_p=0.9,
+                repeat_penalty=1.1,
+                stop=["</s>", "User:", "Assistant:"]
+            )
+            return output["choices"][0]["text"].strip()
         except Exception as e:
             logger.error(f"Geração falhou: {e}")
-            return "Buguei, puto."
+            return "Buguei, kota."
 
 HermesLLM = LocalLLM