Update modules/treinamento.py

modules/treinamento.py

@@ -1,33 +1,33 @@
-# modules/treinamento.py — TREINO LORA ANGOLANO AUTOMÁTICO + EMBEDDINGS (ÉPICO!)
+# modules/treinamento.py — TREINO LEVE (30s) + EMBEDDINGS RÁPIDOS + LORA ANGOLANO
 import json
 import os
 import threading
 import time
 from loguru import logger
 from sentence_transformers import SentenceTransformer
-from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
 from torch.utils.data import Dataset
+import torch
 from .database import Database
 
 # PASTAS
 FINETUNED_PATH = "/home/user/data/finetuned_hermes"
 DATA_PATH = f"{FINETUNED_PATH}/dataset.jsonl"
 EMBEDDINGS_PATH = f"{FINETUNED_PATH}/embeddings.jsonl"
-LORA_PATH = f"{FINETUNED_PATH}/lora_angolano"
+LORA_PATH = f"{FINETUNED_PATH}/lora_leve"
 
 os.makedirs(FINETUNED_PATH, exist_ok=True)
 os.makedirs(LORA_PATH, exist_ok=True)
 
-# MODELOS
-EMBEDDING_MODEL = SentenceTransformer("intfloat/multilingual-e5-large")  # MELHOR PARA PORTUGUÊS/ANGOLANO
-TOKENIZER = None
-MODEL = None
+# EMBEDDING LEVE E RÁPIDO (300MB)
+EMBEDDING_MODEL = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# LOCK PARA THREADS
+# LOCK
 _lock = threading.Lock()
+_dataset = []
 
-class AngolanoDataset(Dataset):
+class LeveDataset(Dataset):
     def __init__(self, data):
         self.data = data
 
@@ -36,31 +36,29 @@ class AngolanoDataset(Dataset):
 
     def __getitem__(self, idx):
         item = self.data[idx]
-        text = f"<s>[INST] {item['user']} [/INST] {item['assistant']}</s>"
-        encoded = TOKENIZER(text, truncation=True, max_length=2048, padding="max_length")
+        text = f"<|im_start|>user\n{item['user']}<|im_end|>\n<|im_start|>assistant\n{item['assistant']}<|im_end|>"
+        encoded = TOKENIZER(text, truncation=True, max_length=512, padding="max_length")
         encoded["labels"] = encoded["input_ids"].copy()
         return {k: torch.tensor(v) for k, v in encoded.items()}
 
-class Treinamento:
+class TreinamentoLeve:
     def __init__(self, db: Database):
         self.db = db
-        self.dataset = []
         self._carregar_dataset()
-        logger.info("TREINAMENTO ANGOLANO ATIVO → EMBEDDINGS + LORA A CADA 4H!")
-
-        # Inicia treino periódico
+        logger.info("TREINAMENTO LEVE ATIVO → LORA A CADA 30 MIN + EMBEDDINGS RÁPIDOS!")
         threading.Thread(target=self._treino_periodico, daemon=True).start()
 
     def _carregar_dataset(self):
+        global _dataset
         if os.path.exists(DATA_PATH):
             with open(DATA_PATH, "r", encoding="utf-8") as f:
-                self.dataset = [json.loads(line) for line in f]
-            logger.info(f"{len(self.dataset)} interações carregadas para treino!")
+                _dataset = [json.loads(line) for line in f]
+            logger.info(f"{len(_dataset)} kandandos carregados!")
 
     def registrar_interacao(self, usuario, mensagem, resposta, numero='', **kwargs):
         self.db.salvar_mensagem(usuario, mensagem, resposta, numero)
         self._salvar_roleplay(mensagem, resposta)
-        self._salvar_embedding(mensagem, resposta)
+        self._salvar_embedding_leve(mensagem, resposta)
 
     def _salvar_roleplay(self, msg, resp):
         entry = {"user": msg, "assistant": resp}
@@ -68,87 +66,78 @@ class Treinamento:
             json.dump(entry, f, ensure_ascii=False)
             f.write("\n")
         with _lock:
-            self.dataset.append(entry)
+            _dataset.append(entry)
 
-    def _salvar_embedding(self, msg, resp):
+    def _salvar_embedding_leve(self, msg, resp):
         try:
-            emb_msg = EMBEDDING_MODEL.encode(msg, normalize_embeddings=True).tolist()
-            emb_resp = EMBEDDING_MODEL.encode(resp, normalize_embeddings=True).tolist()
-            entry = {"msg": msg, "resp": resp, "emb_msg": emb_msg, "emb_resp": emb_resp}
+            emb = EMBEDDING_MODEL.encode(f"{msg} {resp}", normalize_embeddings=True).tolist()
+            entry = {"text": f"{msg} -> {resp}", "emb": emb}
             with open(EMBEDDINGS_PATH, "a", encoding="utf-8") as f:
                 json.dump(entry, f, ensure_ascii=False)
                 f.write("\n")
-            logger.debug("Embedding salvo com sucesso!")
-        except Exception as e:
-            logger.error(f"Erro no embedding: {e}")
+        except: pass
 
     def _treino_periodico(self):
-        global TOKENIZER, MODEL
         while True:
-            time.sleep(4 * 3600)  # A CADA 4 HORAS
-            if len(self.dataset) < 10:
-                logger.info("Poucos dados ainda... esperando mais kandandos!")
+            time.sleep(30 * 60)  # A CADA 30 MINUTOS
+            if len(_dataset) < 5:
                 continue
 
-            logger.info("INICIANDO TREINO LORA ANGOLANO AUTOMÁTICO!")
+            logger.info("INICIANDO TUNE LEVE (30s) → LORA ANGOLANO!")
             try:
-                # Carrega tokenizer e modelo base (só na primeira vez)
-                if TOKENIZER is None:
-                    TOKENIZER = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_fast=True)
+                # Tokenizer (só uma vez)
+                global TOKENIZER
+                if 'TOKENIZER' not in globals():
+                    TOKENIZER = AutoTokenizer.from_pretrained("NousResearch/OpenHermes-2.5-Mistral-7B", use_fast=True)
                     TOKENIZER.pad_token = TOKENIZER.eos_token
 
-                if MODEL is None:
-                    MODEL = AutoModelForCausalLM.from_pretrained(
-                        "meta-llama/Llama-2-7b-hf",
-                        load_in_8bit=True,
-                        device_map="auto"
-                    )
-                    MODEL = prepare_model_for_int8_training(MODEL)
+                # Modelo em 4bit (LEVE!)
+                model = AutoModelForCausalLM.from_pretrained(
+                    "NousResearch/OpenHermes-2.5-Mistral-7B",
+                    load_in_4bit=True,
+                    device_map="auto"
+                )
+                model = prepare_model_for_kbit_training(model)
 
-                # Config LoRA
+                # LoRA leve
                 lora_config = LoraConfig(
-                    r=16,
-                    lora_alpha=32,
+                    r=8,
+                    lora_alpha=16,
                     target_modules=["q_proj", "v_proj"],
                     lora_dropout=0.05,
                     bias="none",
                     task_type="CAUSAL_LM"
                 )
-                model_peft = get_peft_model(MODEL, lora_config)
+                model = get_peft_model(model, lora_config)
 
-                # Dataset
-                train_dataset = AngolanoDataset(self.dataset[-500:])  # últimas 500 interações
+                # Dataset (últimas 50 interações)
+                dataset = LeveDataset(_dataset[-50:])
 
-                # Treino
-                training_args = TrainingArguments(
+                # Treino RÁPIDO
+                args = TrainingArguments(
                     output_dir=LORA_PATH,
-                    per_device_train_batch_size=2,
+                    per_device_train_batch_size=1,
                     gradient_accumulation_steps=4,
-                    num_train_epochs=3,
+                    num_train_epochs=1,
                     learning_rate=2e-4,
                     fp16=True,
-                    logging_steps=10,
-                    save_steps=50,
-                    save_total_limit=2,
+                    logging_steps=5,
+                    save_steps=10,
+                    save_total_limit=1,
                     report_to=[],
-                    disable_tqdm=False
-                )
-
-                trainer = Trainer(
-                    model=model_peft,
-                    args=training_args,
-                    train_dataset=train_dataset,
-                    tokenizer=TOKENIZER
+                    disable_tqdm=True
                 )
 
+                trainer = Trainer(model=model, args=args, train_dataset=dataset)
                 trainer.train()
                 trainer.save_model(LORA_PATH)
-                logger.info("LORA ANGOLANO TREINADO E SALVO! SOTAQUE MAIS FORTE QUE NUNCA!")
 
-                # Converte para GGUF (opcional, se quiser usar direto no llama.cpp)
-                os.system(f"python -m llama_cpp.convert --outfile {FINETUNED_PATH}/adapter_model.gguf --model {LORA_PATH}")
+                # Converte pra GGUF leve
+                os.system(f"python -m llama_cpp.convert --outfile {FINETUNED_PATH}/lora_leve.gguf --model {LORA_PATH} --quantize q4_k_m")
+
+                logger.info("LORA LEVE TREINADO EM 30s! SOTAQUE ANGOLANO + FORTE!")
+                del model, trainer
+                torch.cuda.empty_cache()
 
             except Exception as e:
-                logger.error(f"Erro no treino: {e}")
-                import traceback
-                logger.error(traceback.format_exc())
+                logger.error(f"Erro no tune leve: {e}")