Update train_and_save_models.py

train_and_save_models.py

@@ -10,9 +10,9 @@ from pathlib import Path
 from concurrent.futures import ProcessPoolExecutor, as_completed
 import numpy as np
 import multiprocessing
-import joblib # Importado para salvar e carregar modelos
+import joblib # Importado para salvar os modelos
 
-# REMOVIDOS Flask e CORS, pois este script é apenas para treinamento local.
+# REMOVIDAS importações Flask e CORS, pois este script é apenas para treinamento local.
 # from flask import Flask, request, jsonify
 # from flask_cors import CORS
 
@@ -46,11 +46,13 @@ torch.manual_seed(RANDOM_SEED)
 if torch.cuda.is_available():
     torch.cuda.manual_seed_all(RANDOM_SEED)
 
-model_base = None
+model_base = None # Será inicializado durante o treinamento
 mlp_regressor, scaler = None, None
 tfidf_vectorizer, tfidf_regressor = None, None
 
-# --- Vocabulário Expandido (compartilhado) ---
+# --- Vocabulário Expandido (compartilhado para geração de dados) ---
+# Estas listas são usadas SOMENTE no trainer.py para gerar os dados de treinamento.
+# Elas serão duplicadas (ou uma versão reduzida) no app.py se o "Gerar Evento Aleatório" for mantido.
 ADVERSARIAL_RISK_ACTORS = [
     "Unsandboxed process", "Leaked API key", "Misconfigured service account", "Shadow IT application", 
     "Dormant user account", "Ransomware payload", "Phishing attempt", "Insider threat", 
@@ -237,23 +239,23 @@ ADVERSARIAL_SAFE_TARGETS = [
     "Kubernetes cluster security posture", "Docker container security configuration", "AWS cloud infrastructure",
     "Azure cloud resources", "GCP cloud services", "container orchestration security",
     "serverless function security", "cloud API security", "microservice security architecture",
-    "container registry security", "cloud logging security", "infrastructure as code security",
+    "container registry security", "cloud logging security", "infrastructure as code repository",
     "Git repository security", "CI/CD pipeline security", "Docker image security",
     "artifact repository security", "infrastructure provisioning security", "secret management vault",
     "code signing certificate store", "dependency management system", "deployment automation platform",
-    "build environment security", "CI/CD security scanning", "infrastructure monitoring security",
+    "build environment with elevated privileges", "CI/CD security scanning tools", "infrastructure monitoring system",
     "industrial control system security", "SCADA system security", "IoT device security",
     "edge computing security", "smart city infrastructure security", "medical device network",
     "automotive system security", "home automation security", "sensor security",
-    "industrial protocol security", "edge gateway security", "IoT device management security",
-    "enterprise mobile device security", "mobile app security", "mobile device management security",
-    "mobile banking security", "mobile certificate security", "mobile security scanning",
-    "BYOD policy security", "mobile endpoint security", "mobile app store security",
-    "mobile device fingerprinting security", "mobile phishing protection", "mobile security testing",
-    "network segmentation security", "firewall security", "VPN security",
-    "DNS security", "BGP routing security", "wireless network security",
-    "Bluetooth security", "NFC security", "network monitoring security",
-    "traffic analysis security", "protocol security", "network infrastructure security"
+    "industrial protocol gateway", "edge security monitoring system", "IoT device firmware repository",
+    "enterprise mobile device fleet", "mobile app store backend", "mobile device management system",
+    "mobile banking infrastructure", "mobile certificate authority", "mobile security scanning service",
+    "BYOD policy enforcement system", "mobile endpoint detection system", "mobile app security testing platform",
+    "mobile device fingerprinting database", "mobile phishing detection system", "mobile app code signing service",
+    "network segmentation firewall", "VPN concentrator", "DNS authoritative server",
+    "BGP route reflector", "wireless access point controller", "network monitoring system",
+    "traffic analysis platform", "network security scanning tool", "protocol analysis system",
+    "network infrastructure management", "security information system", "network forensics platform"
 ]
 ADVERSARIAL_SAFE_OUTCOMES = [
     "all tests passed, security posture confirmed", "the configuration was hardened as per policy", 
@@ -299,6 +301,7 @@ LOW_RISK_KEYWORDS = {
     'backup completed': -20, 'schema migration successful': -15, 'network policy updated': -10
 }
 
+# --- Funções de Geração de Dados de TREINAMENTO (Base Sólida) ---
 def generate_event_text_for_training(is_risk: bool) -> tuple[str, float]:
     if is_risk:
         actor = random.choice(ADVERSARIAL_RISK_ACTORS)
@@ -351,8 +354,24 @@ def populate_database_initial():
     conn.close()
     print("Banco de dados populado inicialmente com sucesso.")
 
+# --- Funções de Embedding e Treinamento ---
+def init_sbert_worker():
+    global model_base
+    if model_base is None:
+        print(f"Processo worker {os.getpid()} carregando o modelo {MODEL_NAME}...")
+        model_base = SentenceTransformer(MODEL_NAME)
+        torch.set_num_threads(1)
+
+def extract_embeddings_batch_worker(texts: list[str]) -> list[list[float]]:
+    global model_base
+    if model_base is None:
+        raise RuntimeError("SentenceTransformer não foi inicializado no worker.")
+        
+    embeddings = model_base.encode(texts, convert_to_numpy=True, show_progress_bar=False)
+    return embeddings.tolist()
+
 def train_and_save_all_models():
-    global mlp_regressor, scaler, tfidf_vectorizer, tfidf_regressor, model_base
+    global mlp_regressor, scaler, tfidf_vectorizer, tfidf_regressor # Model_base é para os workers, não para o principal
     print("Iniciando o treinamento de todos os modelos a partir do banco de dados...")
     
     conn = sqlite3.connect(DB_NAME)
@@ -367,6 +386,7 @@ def train_and_save_all_models():
     train_texts = [row[0] for row in train_data]
     y_train = np.array([row[1] for row in train_data])
     
+    # --- Cabeça 1: Embedding Profundo (MLPRegressor) ---
     print("1. Treinando modelo de Embedding Profundo (MLPRegressor)...")
     X_train_embeddings = []
     
@@ -406,6 +426,7 @@ def train_and_save_all_models():
     
     print("   ... modelo de Embedding Profundo treinado.")
 
+    # --- Cabeça 2: Vetorial Clássico (TF-IDF) ---
     print("2. Treinando modelo Vetorial Clássico (TF-IDF + Ridge)...")
     tfidf_vectorizer = TfidfVectorizer(ngram_range=(1, 3), min_df=5, max_df=0.7, max_features=10000) 
     X_train_tfidf = tfidf_vectorizer.fit_transform(train_texts)