""" Implementação concreta (Keras/TensorFlow) do modelo treinado pelo pipeline. Arquitetura: MLP simples para dados tabulares. A configuração de camadas, dropout, otimizador, etc. vem de ``config`` (lido do ``config.json`` do treino) ou de defaults razoáveis. A classe ``TabularMLP`` herda de ``BaseModel`` (em ``docker/base_model.py``) para manter a API uniforme entre frameworks. Uma função ``create_model`` fábrica é exposta no fim do arquivo para que ``main.py`` continue funcionando sem mudanças. """ from __future__ import annotations import os import sys from typing import Any, Dict, List, Optional, Tuple import numpy as np import tensorflow as tf from tensorflow import keras from tensorflow.keras import layers # Permite importar base_model.py de ../ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from base_model import BaseModel # noqa: E402 from utils import emit_progress # noqa: E402 # --------------------------------------------------------------------------- # Wrapper que implementa a API de BaseModel para Keras # --------------------------------------------------------------------------- class TabularMLP(BaseModel): """MLP tabular em Keras/TensorFlow, plugável no pipeline.""" def build_model(self) -> keras.Model: hp = self.config.get("hyperparameters", {}) hidden_dims = self.config.get("hidden_dims") or hp.get("hidden_dims") or [ self.config.get("hidden_dim1", 32), self.config.get("hidden_dim2", 16), ] dropout = float(hp.get("dropout", self.config.get("dropout", 0.0))) # Constrói o modelo Sequential model = keras.Sequential() model.add(layers.InputLayer(input_shape=(self.input_size,))) # Adiciona camadas ocultas for h in hidden_dims: model.add(layers.Dense(h, activation="relu")) if dropout and dropout > 0: model.add(layers.Dropout(dropout)) # Camada de saída model.add(layers.Dense(self.num_classes, activation="softmax")) self._hidden_dims = list(hidden_dims) self._dropout = dropout return model # ---------------- treinamento ---------------- def train_model( self, train_loader: Any, val_loader: Optional[Any], epochs: int, lr: float, **kwargs: Any, ) -> Dict[str, list]: hp = self.config.get("hyperparameters", {}) epochs = int(hp.get("epochs", epochs)) lr = float(hp.get("learning_rate", lr)) weight_decay = float(hp.get("weight_decay", 0.0)) optimizer_name = str(hp.get("optimizer", "adam")).lower() # Desempacota tuplas (X, y) do dataset if isinstance(train_loader, tuple) and len(train_loader) == 2: X_train, y_train = train_loader else: raise ValueError(f"train_loader deve ser tupla (X, y), recebeu {type(train_loader)}") if isinstance(val_loader, tuple) and len(val_loader) == 2: X_val, y_val = val_loader else: raise ValueError(f"val_loader deve ser tupla (X, y), recebeu {type(val_loader)}") # Compila o modelo optimizer = self._build_optimizer(optimizer_name, lr, weight_decay) self.model.compile( optimizer=optimizer, loss="sparse_categorical_crossentropy", metrics=["accuracy"], ) history = { "train_loss": [], "train_acc": [], "val_loss": [], "val_acc": [], } n_batches = max((len(X_train) + 31) // 32, 1) # Batches por epoch total_train_steps = max(epochs * n_batches, 1) batch_size = int(hp.get("batch_size", 32)) best_val_acc = -float("inf") best_epoch = -1 for epoch in range(1, epochs + 1): # Treina 1 época hist = self.model.fit( X_train, y_train, batch_size=batch_size, epochs=1, validation_data=(X_val, y_val), verbose=0, ) # Extrai histórico train_loss = float(hist.history["loss"][0]) train_acc = float(hist.history["accuracy"][0]) val_loss = float(hist.history["val_loss"][0]) val_acc = float(hist.history["val_accuracy"][0]) history["train_loss"].append(round(train_loss, 6)) history["train_acc"].append(round(train_acc, 6)) history["val_loss"].append(round(val_loss, 6)) history["val_acc"].append(round(val_acc, 6)) # Emite progresso por batch (simula progresso contínuo por época) for batch_idx in range(1, n_batches + 1): global_step = (epoch - 1) * n_batches + batch_idx inner_pct = int(global_step * 100 / total_train_steps) emit_progress(inner_pct, 1) # Rastreia melhor validação if val_acc > best_val_acc: best_val_acc = val_acc best_epoch = epoch print( f"Epoch [{epoch:>3}/{epochs}] " f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} " f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}" ) history["best_epoch"] = best_epoch history["best_val_acc"] = round(best_val_acc, 6) if best_epoch > 0 else None self.history = history return history # ---------------- avaliação ---------------- def evaluate(self, data_loader: Any) -> Dict[str, float]: if isinstance(data_loader, tuple) and len(data_loader) == 2: X_val, y_val = data_loader else: X_val, y_val = data_loader[2], data_loader[3] loss, accuracy = self.model.evaluate(X_val, y_val, verbose=0) return { "loss": float(loss), "accuracy": float(accuracy), "n": len(X_val), } # ---------------- inferência ---------------- def predict(self, inputs: Any) -> np.ndarray: if isinstance(inputs, tuple): # Se for tuple de (X_train, y_train, X_val, y_val), usa X_val X = inputs[2].astype(np.float32) elif isinstance(inputs, np.ndarray): X = inputs.astype(np.float32) else: X = np.array(inputs, dtype=np.float32) if X.ndim == 1: X = X.reshape(1, -1) # Predições como classe logits = self.model.predict(X, verbose=0) return np.argmax(logits, axis=1) # ---------------- persistência ---------------- def save_model(self, filename: str) -> str: path = os.path.join(self.model_dir, filename) # Remove extensão se for .pt (compatibilidade) if path.endswith(".pt"): path = path[:-3] + ".keras" self.model.save(path) print(f"✓ Model saved to {path}") return path def load_model(self, filename: str) -> None: path = filename if os.path.isabs(filename) else os.path.join( self.model_dir, filename ) # Converte .pt para .keras if path.endswith(".pt"): path = path[:-3] + ".keras" self.model = keras.models.load_model(path) # ---------------- helpers ---------------- def _build_optimizer( self, name: str, lr: float, weight_decay: float ) -> keras.optimizers.Optimizer: if name == "sgd": return keras.optimizers.SGD(learning_rate=lr, weight_decay=weight_decay) if name == "rmsprop": return keras.optimizers.RMSprop(learning_rate=lr, weight_decay=weight_decay) # default: adam return keras.optimizers.Adam(learning_rate=lr, weight_decay=weight_decay) # --------------------------------------------------------------------------- # Fábrica usada pelo main.py # --------------------------------------------------------------------------- def create_model( input_size: int, num_classes: int, project_name: str, base_path: str, config: Optional[Dict[str, Any]] = None, ) -> TabularMLP: """Cria um ``TabularMLP`` pronto pra uso pelo main.py.""" return TabularMLP( input_size=input_size, num_classes=num_classes, project_name=project_name, base_path=base_path, framework="keras", config=config or {}, )