Update modules/treinamento.py

modules/treinamento.py

@@ -5,26 +5,21 @@ import json
 import os
 from loguru import logger
 from sentence_transformers import SentenceTransformer
-from transformers import TrainingArguments, Trainer
-from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
+from llama_cpp import Llama
 from .database import Database
+from .local_llm import _get_llm  # ← REUSO DO MESMO LLAMA OBJECT
 
 # === CONFIGURAÇÕES GLOBAIS ===
 EMBEDDING_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"
 embedding_model = SentenceTransformer(EMBEDDING_MODEL)
 
-HERMES_PATH = "/app/models/hermes-7b"
 FINETUNED_PATH = "/app/data/finetuned_hermes"
 os.makedirs(FINETUNED_PATH, exist_ok=True)
 
-# === IMPORTA SINGLETON DO local_llm.py (MESMO OBJETO!) ===
-from .local_llm import _get_hermes_singleton  # ← REUSO 100%
-
-
 def gerar_embedding(text: str):
     return embedding_model.encode(text, convert_to_numpy=True)
 
-
 class Treinamento:
     def __init__(self, db: Database, interval_hours: int = 4):
         self.db = db
@@ -32,38 +27,25 @@ class Treinamento:
         self._thread = None
         self._running = False
 
-        # REUSO DO MODELO JÁ CARREGADO PELO local_llm.py
-        result = _get_hermes_singleton()
-        if result:
-            self.base_model, self.tokenizer = result  # ← JÁ É (model, tokenizer)
-            # Prepara LoRA (sem recarregar!)
-            self.model = prepare_model_for_kbit_training(self.base_model)
-            peft_config = LoraConfig(
-                r=32,
-                lora_alpha=64,
-                target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-                lora_dropout=0.05,
-                bias="none",
-                task_type="CAUSAL_LM"
-            )
-            self.model = get_peft_model(self.model, peft_config)
-            logger.info("Treinamento conectado ao Hermes 7B GLOBAL (sem reload!)")
+        # REUSO DO LLAMA JÁ CARREGADO
+        self.llm = _get_llm()
+        if self.llm:
+            logger.info("Treinamento conectado ao OpenHermes GGUF GLOBAL (sem reload!)")
         else:
-            self.base_model = self.tokenizer = self.model = None
-            logger.warning("Hermes não disponível para finetune")
+            logger.warning("LLM não disponível para finetune")
 
     def registrar_interacao(self, usuario, mensagem, resposta, numero='', is_reply=False, mensagem_original=''):
         self.db.salvar_mensagem(usuario, mensagem, resposta, numero, is_reply, mensagem_original)
         self._aprender_roleplay(numero, mensagem, resposta)
 
     def _aprender_roleplay(self, numero: str, msg: str, resp: str):
-        if not numero or not self.model:
+        if not numero or not self.llm:
             return
 
         dataset_path = f"{FINETUNED_PATH}/roleplay.jsonl"
         entry = {
             "messages": [
-                {"role": "system", "content": "Tu és Akira, angolana sarcástica de Luanda. Fala com gírias: bué, fixe, kota, mwangolé, kandando."},
+                {"role": "system", "content": "Tu és Akira, angolana sarcástica de Luanda. Fala com gírias: bué, fixe, kota, mwangolé, kandando, na boa."},
                 {"role": "user", "content": msg},
                 {"role": "assistant", "content": resp}
             ]
@@ -74,7 +56,7 @@ class Treinamento:
         logger.debug(f"Roleplay salvo: {msg[:30]}... → {resp[:30]}...")
 
     def train_once(self):
-        if not self.model:
+        if not self.llm:
             logger.warning("Modelo não carregado. Pulando finetune.")
             return
 
@@ -83,65 +65,59 @@ class Treinamento:
             logger.info("Poucos dados. Esperando mais interações...")
             return
 
-        logger.info("INICIANDO FINETUNE LORA (ROLEPLAY ANGOLANO)...")
+        logger.info("INICIANDO FINETUNE LORA (ANGOLANO STYLE) COM llama.cpp...")
 
         try:
-            from datasets import load_dataset
-
-            dataset = load_dataset("json", data_files=dataset_path, split="train")
-
-            def tokenize_function(examples):
-                # Junta mensagens em um texto só
-                texts = []
-                for msgs in examples["messages"]:
-                    text = ""
-                    for msg in msgs:
-                        if msg["role"] == "system":
-                            text += f"<|system|>{msg['content']}<|end|>"
-                        elif msg["role"] == "user":
-                            text += f"<|user|>{msg['content']}<|end|>"
-                        elif msg["role"] == "assistant":
-                            text += f"<|assistant|>{msg['content']}<|endoftext|>"
-                    texts.append(text)
-                return self.tokenizer(texts, truncation=True, max_length=512, padding="max_length")
-
-            tokenized = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
-
-            training_args = TrainingArguments(
-                output_dir=FINETUNED_PATH,
-                per_device_train_batch_size=1,
-                gradient_accumulation_steps=4,
-                num_train_epochs=1,
-                learning_rate=2e-4,
-                fp16=True,
-                logging_steps=5,
-                save_steps=50,
-                save_total_limit=2,
-                report_to=[],
-                disable_tqdm=False,
-                dataloader_num_workers=0,
-                remove_unused_columns=False
-            )
-
-            trainer = Trainer(
-                model=self.model,
-                args=training_args,
-                train_dataset=tokenized
+            # === CRIA CONFIG LORA ===
+            lora_config = LoraConfig(
+                r=32,
+                lora_alpha=64,
+                target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+                lora_dropout=0.05,
+                bias="none",
+                task_type="CAUSAL_LM"
             )
 
-            trainer.train()
-
-            # SALVA APENAS O LORA
-            self.model.save_pretrained(FINETUNED_PATH)
-            self.tokenizer.save_pretrained(FINETUNED_PATH)
-            logger.info("FINETUNE CONCLUÍDO! LORA ANGOLANO SALVO!")
+            # === APLICA LORA NO LLAMA (cria arquivo temporário) ===
+            lora_path = f"{FINETUNED_PATH}/temp_lora"
+            os.makedirs(lora_path, exist_ok=True)
+
+            # Salva config
+            lora_config.save_pretrained(lora_path)
+
+            # === TREINA COM llama.cpp CLI (mais rápido e estável) ===
+            cmd = [
+                "python", "-m", "llama_cpp.convert",
+                "--outfile", f"{lora_path}/adapter_model.bin",
+                "--model", "/app/models/openhermes-2.5-mistral-7b.Q4_K_M.gguf",
+                "--lora-out", lora_path,
+                "--train", dataset_path,
+                "--epochs", "1",
+                "--lora-r", "32",
+                "--lora-alpha", "64",
+                "--batch", "4",
+                "--threads", "4",
+                "--ctx", "4096"
+            ]
 
-            # Limpa dataset
-            open(dataset_path, 'w').close()
-            logger.info("Dataset limpo. Pronto pro próximo ciclo.")
+            import subprocess
+            result = subprocess.run(cmd, capture_output=True, text=True)
+            if result.returncode == 0:
+                logger.info("FINETUNE LORA CONCLUÍDO COM SUCESSO!")
+                # Move pro lugar certo
+                import shutil
+                if os.path.exists(f"{lora_path}/adapter_model.bin"):
+                    shutil.move(f"{lora_path}/adapter_model.bin", f"{FINETUNED_PATH}/adapter_model.bin")
+                    shutil.move(f"{lora_path}/adapter_config.json", f"{FINETUNED_PATH}/adapter_config.json")
+                    logger.info("LORA ANGOLANO SALVO EM /app/data/finetuned_hermes")
+                # Limpa dataset
+                open(dataset_path, 'w').close()
+                logger.info("Dataset limpo. Pronto pro próximo ciclo.")
+            else:
+                logger.error(f"Erro no treino: {result.stderr}")
 
         except Exception as e:
-            logger.error(f"Erro no finetune: {e}")
+            logger.error(f"Erro crítico no finetune: {e}")
             import traceback
             logger.error(traceback.format_exc())
 
@@ -155,9 +131,9 @@ class Treinamento:
             time.sleep(interval)
 
     def start_periodic_training(self):
-        if self._running or not self.model:
+        if self._running or not self.llm:
             return
         self._running = True
         self._thread = threading.Thread(target=self._run_loop, daemon=True)
         self._thread.start()
-        logger.info(f"Treinamento periódico iniciado (a cada {self.interval_hours}h)")
+        logger.info(f"Treinamento periódico iniciado (a cada {self.interval_hours}h)")