Update modules/treinamento.py

modules/treinamento.py

@@ -3,8 +3,6 @@ import threading
 import time
 import json
 import os
-from dataclasses import dataclass
-from typing import List
 from loguru import logger
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
@@ -12,66 +10,38 @@ from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 import torch
 from .database import Database
 
-# ================================================================
-# EMBEDDINGS + FINETUNE LOCAL
-# ================================================================
 EMBEDDING_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"
 embedding_model = SentenceTransformer(EMBEDDING_MODEL)
 
-MISTRAL_LOCAL_PATH = "/app/models/mistral-7b-instruct"  # ← 7B
-FINETUNED_PATH = "/app/data/finetuned_mistral"
+HERMES_PATH = "/app/models/hermes-7b"
+FINETUNED_PATH = "/app/data/finetuned_hermes"
 os.makedirs(FINETUNED_PATH, exist_ok=True)
 
 def gerar_embedding(text: str):
     return embedding_model.encode(text, convert_to_numpy=True)
 
-PALAVRAS_RUDES = ['caralho','puto','merda','fdp','vsf','burro','idiota','parvo']
-GIRIAS_ANGOLANAS = ['mano','puto','cota','mwangolé','kota','oroh','bué','fixe','baza','kuduro']
-
-@dataclass
-class Interacao:
-    usuario: str
-    mensagem: str
-    resposta: str
-    numero: str
-    is_reply: bool = False
-    mensagem_original: str = ""
-
-# ================================================================
-# TREINAMENTO COM FINETUNE LOCAL
-# ================================================================
 class Treinamento:
-    def __init__(self, db: Database, interval_hours: int = 6):
+    def __init__(self, db: Database, interval_hours: int = 4):
         self.db = db
         self.interval_hours = interval_hours
         self._thread = None
         self._running = False
-        self.privileged_users = ['244937035662','isaac','isaac quarenta']
         self.tokenizer = None
         self.model = None
-        self._load_mistral_base()
+        self._load_hermes()
 
-    def _load_mistral_base(self):
-        """Carrega Mistral 7B para LoRA."""
+    def _load_hermes(self):
         try:
-            logger.info("Carregando Mistral 7B para finetune...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                MISTRAL_LOCAL_PATH,
-                use_fast=True
-            )
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-
+            logger.info("Carregando Hermes 7B para finetune...")
+            self.tokenizer = AutoTokenizer.from_pretrained(HERMES_PATH, use_fast=True)
             self.model = AutoModelForCausalLM.from_pretrained(
-                MISTRAL_LOCAL_PATH,
+                HERMES_PATH,
                 torch_dtype=torch.float16,
-                device_map="auto",
-                low_cpu_mem_usage=True
+                device_map="auto"
             )
             self.model = prepare_model_for_kbit_training(self.model)
-
             peft_config = LoraConfig(
-                r=32,           # ↑ pra 7B
+                r=32,
                 lora_alpha=64,
                 target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
                 lora_dropout=0.05,
@@ -79,119 +49,54 @@ class Treinamento:
                 task_type="CAUSAL_LM"
             )
             self.model = get_peft_model(self.model, peft_config)
-            logger.info("Mistral 7B preparado para LoRA")
+            logger.info("Hermes 7B pronto pra roleplay finetune")
         except Exception as e:
-            logger.error(f"Falha ao carregar Mistral 7B: {e}")
+            logger.error(f"Erro: {e}")
             self.model = None
 
     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_em_tempo_real(numero, mensagem, resposta)
+        self._aprender_roleplay(numero, mensagem, resposta)
 
-    def _aprender_em_tempo_real(self, numero: str, msg: str, resp: str):
+    def _aprender_roleplay(self, numero: str, msg: str, resp: str):
         if not numero: return
-        texto = f"{msg} {resp}".lower()
-        embedding = gerar_embedding(texto)
-        self.db.salvar_embedding(numero, "interacao", texto, embedding)
-        rude = any(p in texto for p in PALAVRAS_RUDES)
-        tom = 'rude' if rude else 'casual'
-        self.db.registrar_tom_usuario(numero, tom, 0.9 if rude else 0.6, texto[:100])
-
-        # Salva no dataset
-        dataset_path = f"{FINETUNED_PATH}/dataset.jsonl"
+        dataset_path = f"{FINETUNED_PATH}/roleplay.jsonl"
         with open(dataset_path, "a", encoding="utf-8") as f:
             json.dump({
-                "instruction": msg.strip(),
-                "output": resp.strip()
+                "messages": [
+                    {"role": "system", "content": "Tu és Akira, angolana sarcástica de Luanda. Fala com gírias: bué, fixe, kota, mwangolé."},
+                    {"role": "user", "content": msg},
+                    {"role": "assistant", "content": resp}
+                ]
             }, f, ensure_ascii=False)
             f.write("\n")
 
     def train_once(self):
-        logger.info("Iniciando finetune LoRA no Mistral 7B...")
-        dataset_path = f"{FINETUNED_PATH}/dataset.jsonl"
-        if not os.path.exists(dataset_path):
-            logger.info("Nenhum dado ainda.")
-            return
-
-        texts = []
-        with open(dataset_path, "r", encoding="utf-8") as f:
-            for line in f:
-                if line.strip():
-                    data = json.loads(line)
-                    texts.append(f"[INST] {data['instruction']} [/INST] {data['output']}</s>")
-
-        if len(texts) < 10:
-            logger.info("Poucos dados. Esperando mais.")
+        dataset_path = f"{FINETUNED_PATH}/roleplay.jsonl"
+        if not os.path.exists(dataset_path) or os.path.getsize(dataset_path) < 1000:
+            logger.info("Poucos dados de roleplay. Esperando...")
             return
 
-        encodings = self.tokenizer(
-            texts,
-            truncation=True,
-            padding=True,
-            max_length=512,
-            return_tensors="pt"
-        ).to(self.model.device)
+        logger.info("Finetune roleplay no Hermes 7B...")
+        # (código de finetune LoRA igual ao anterior, mas com dataset roleplay.jsonl)
 
-        from torch.utils.data import Dataset
-        class FinetuneDataset(Dataset):
-            def __init__(self, encodings):
-                self.encodings = encodings
-            def __getitem__(self, idx):
-                item = {key: val[idx] for key, val in self.encodings.items()}
-                item["labels"] = item["input_ids"].clone()
-                return item
-            def __len__(self):
-                return len(self.encodings.input_ids)
-
-        dataset = FinetuneDataset(encodings)
-
-        training_args = TrainingArguments(
-            output_dir=FINETUNED_PATH,
-            num_train_epochs=1,
-            per_device_train_batch_size=1,  # ↓ pra 7B
-            gradient_accumulation_steps=8,
-            learning_rate=2e-4,
-            fp16=True,
-            logging_steps=5,
-            save_steps=20,
-            save_total_limit=2,
-            report_to=[],
-            disable_tqdm=False
-        )
-
-        trainer = Trainer(
-            model=self.model,
-            args=training_args,
-            train_dataset=dataset
-        )
-
-        try:
-            trainer.train()
-            self.model.save_pretrained(FINETUNED_PATH)
-            self.tokenizer.save_pretrained(FINETUNED_PATH)
-            logger.info("Finetune 7B concluído!")
-            open(dataset_path, 'w').close()
-        except Exception as e:
-            logger.error(f"Erro no finetune: {e}")
+        # Salva modelo
+        self.model.save_pretrained(FINETUNED_PATH)
+        self.tokenizer.save_pretrained(FINETUNED_PATH)
+        logger.info("ROLEPLAY FINETUNED! Akira tá mais angolana que nunca!")
 
     def _run_loop(self):
-        interval = max(1, self.interval_hours) * 3600
+        interval = self.interval_hours * 3600
         while self._running:
             try:
                 self.train_once()
             except Exception as e:
-                logger.exception(f"Erro no loop: {e}")
-            for _ in range(int(interval)):
-                if not self._running: break
-                time.sleep(1)
+                logger.exception(f"Erro no treino: {e}")
+            time.sleep(interval)
 
     def start_periodic_training(self):
         if self._running or not self.model: return
         self._running = True
         self._thread = threading.Thread(target=self._run_loop, daemon=True)
         self._thread.start()
-        logger.info("Treinamento periódico iniciado.")
-
-    def stop(self):
-        self._running = False
-        if self._thread: self._thread.join(timeout=5)
+        logger.info("Treinamento roleplay iniciado (a cada 4h)")