Merge pull request #9 from Aignosi/feature/SIENTIAPDE-1321

Feature/SIENTIAPDE-1321: A equação do modelo deve ficar disponível como artefato logo após o treinamento.

.gitignore

@@ -145,6 +145,7 @@ celerybeat.pid
 env/
 venv/
 ENV/
+venv_311/
 env.bak/
 venv.bak/
 

doc/model_equation_implementation.md

@@ -0,0 +1,78 @@
+# Implementação da Equação do Modelo como Artefato JSON
+
+## Visão Geral
+
+Esta implementação adiciona a capacidade de extrair e salvar a equação do modelo de regressão linear como um artefato JSON, seguindo a arquitetura existente do projeto.
+
+## Mudanças Implementadas
+
+### 1. TrainModelResult
+- **Arquivo**: `model_manager/utils/models/train_model_result.py`
+- **Mudanças**:
+  - Adicionado campo `equation: dict | None = None` para armazenar os metadados da equação
+  - Adicionado campo `equation_path: str | None = None` para armazenar o caminho do arquivo JSON
+
+### 2. TrainingRepository
+- **Arquivo**: `model_manager/utils/repository/training_repository.py`
+- **Mudanças**:
+  - Adicionado método `_extract_model_equation()` para extrair coeficientes e intercept do modelo
+  - Integrado a extração da equação no método `after_train_calculation()`
+
+### 3. ModelRepository
+- **Arquivo**: `model_manager/utils/repository/model_repository.py`
+- **Mudanças**:
+  - Adicionado import do módulo `json`
+  - Modificado `_generate_report()` para salvar a equação como arquivo JSON
+  - Modificado `_save_run()` para fazer log do artefato da equação no MLflow
+
+## Estrutura do JSON da Equação
+
+O arquivo `model_equation.json` terá a seguinte estrutura:
+
+```json
+{
+  "target_variable": "target_column_name",
+  "coefficients": {
+    "feature1": 0.123456,
+    "feature2": -0.789012,
+    "feature3": 0.345678
+  },
+  "intercept": 1.234567,
+  "equation_string": "target_column_name = 1.234567 + 0.123456 * feature1 + -0.789012 * feature2 + 0.345678 * feature3",
+  "latex_equation": "target_column_name = 1.234567 + 0.123456 \\cdot feature1 + -0.789012 \\cdot feature2 + 0.345678 \\cdot feature3",
+  "model_type": "Linear Regression"
+}
+```
+
+## Fluxo de Execução
+
+1. **Treinamento**: O modelo é treinado no `TrainingRepository.train()`
+2. **Pós-treinamento**: O método `after_train_calculation()` é chamado, que:
+   - Calcula as métricas (MSE, MAE, R²)
+   - Extrai a equação usando `_extract_model_equation()`
+3. **Salvamento**: O `ModelRepository.save_model()` é chamado, que:
+   - Gera os artefatos (relatórios, dados CSV)
+   - Salva a equação como `model_equation.json`
+   - Faz log de todos os artefatos no MLflow
+
+## Benefícios
+
+- **Rastreabilidade**: A equação fica disponível como artefato versionado no MLflow
+- **Transparência**: Fácil acesso aos coeficientes e estrutura do modelo
+- **Compatibilidade**: Formato JSON facilita integração com outras ferramentas
+- **Flexibilidade**: Inclui tanto formato legível quanto LaTeX para diferentes usos
+
+## Compatibilidade
+
+Esta implementação é totalmente compatível com:
+- A arquitetura existente do projeto
+- O fluxo de treinamento atual
+- O sistema de logging do MLflow
+- Os testes existentes (não quebra funcionalidades)
+
+## Exemplo de Uso
+
+Após o treinamento, a equação estará disponível em:
+- **Memória**: `train_result.equation` (dicionário Python)
+- **Arquivo**: `train_result.equation_path` (caminho para o JSON)
+- **MLflow**: Como artefato `model_equation.json` no run do experimento

model_manager/utils/models/train_model_result.py

@@ -28,6 +28,8 @@ class TrainModelResult:
         mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
         mae_val (float | None): The Mean Absolute Error (MAE) of the predictions. Default is None.
         r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
+        equation (dict | None): The equation of the model. Default is None.
+        equation_path (str | None): The path to the equation file. Default is None.
         run_name (str | None): The name of the MLFlow run. Default is None.
         report_path (str | None): The path to the generated HTML report file. Default is None.
         train_data_path (str | None): The path to the training dataset CSV file. Default is None.
@@ -47,6 +49,8 @@ class TrainModelResult:
     mse_val: float | None = None
     mae_val: float | None = None
     r2_val: float | None = None
+    equation: dict | None = None
+    equation_path: str | None = None
     run_name: str | None = None
     report_path: str | None = None
     train_data_path: str | None = None

model_manager/utils/repository/model_repository.py

@@ -9,6 +9,7 @@ and logging model runs to MLFlow.
 
 """
 
+import json
 import os
 import shutil
 import warnings
@@ -224,6 +225,10 @@ class ModelRepository:
             self.model_serving.log_artifact(data.train_data_path)
             self.model_serving.log_artifact(data.test_data_path)
 
+            # Log equation artifact if available
+            if data.equation_path and path.exists(data.equation_path):
+                self.model_serving.log_artifact(data.equation_path)
+
     def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
         """
         Prepares the reference and current datasets for artifact generation.
@@ -412,6 +417,12 @@ class ModelRepository:
             data.test_data_path = path.join(data.run_dir, 'test_data.csv')
             current_data.to_csv(data.test_data_path, index=False)
 
+            # Save equation as JSON
+            if data.equation is not None:
+                data.equation_path = path.join(data.run_dir, 'model_equation.json')
+                with open(data.equation_path, 'w', encoding='utf-8') as f:
+                    json.dump(data.equation, f, indent=2, ensure_ascii=False)
+
             return data
         except ValueError as e:
             error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'

model_manager/utils/repository/training_repository.py

@@ -175,6 +175,10 @@ class TrainingRepository:
         )
 
         tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
+
+        # Extract model equation
+        tmr.equation = self._extract_model_equation(tmr.regr, params)
+
         self.logger.info(
             f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
         )
@@ -249,3 +253,52 @@ class TrainingRepository:
             scaler_params={} if params.use_scaler else None,
             ar_var=params.target_variable if params.include_ar else None,
         )
+
+    def _extract_model_equation(
+        self, regr: LinearRegressionModel, params: TrainModelParams
+    ) -> dict:
+        """
+        Extract the linear regression equation coefficients and create equation metadata.
+
+        This method extracts the coefficients and intercept from the trained model
+        and creates a structured dictionary containing the equation information
+        for serialization as JSON artifact.
+
+        Args:
+            regr: Trained LinearRegressionModel object
+            params: Training parameters containing variable information
+
+        Returns:
+            dict: Equation metadata containing:
+                - target_variable: Name of the target variable
+                - coefficients: Dictionary mapping variable names to coefficients
+                - intercept: Model intercept value
+                - equation_string: Human-readable equation string
+                - latex_equation: LaTeX formatted equation
+        """
+        coefficients = regr.regr.coef_
+        intercept = regr.regr.intercept_
+
+        # Create coefficients dictionary
+        coefficients_dict = {}
+        for i, var in enumerate(params.variable_columns):
+            coefficients_dict[var] = float(coefficients[i])
+
+        # Create equation string
+        equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
+        equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(
+            equation_parts
+        )
+
+        # Create LaTeX equation
+        latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()]
+        latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts)
+
+        return {
+            'target_variable': params.target_variable,
+            'coefficients': coefficients_dict,
+            'intercept': float(intercept),
+            'equation_string': equation_string,
+            'latex_equation': latex_equation,
+            'model_type': 'Linear Regression',
+        }

tests/utils/models/test_train_model_result.py

@@ -175,11 +175,11 @@ def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
 
 
 def test_train_model_result_field_count():
-    """Test that TrainModelResult has exactly 17 fields."""
+    """Test that TrainModelResult has exactly 19 fields."""
     from dataclasses import fields
 
     result_fields = fields(TrainModelResult)
-    assert len(result_fields) == 17
+    assert len(result_fields) == 19
 
     field_names = {f.name for f in result_fields}
     expected_fields = {
@@ -195,6 +195,8 @@ def test_train_model_result_field_count():
         'mse_val',
         'mae_val',
         'r2_val',
+        'equation',
+        'equation_path',
         'run_name',
         'report_path',
         'train_data_path',

tests/utils/repository/test_model_repository.py

@@ -197,5 +197,5 @@ def test_get_reports_directory(mock_model_serving_class, mock_logger):
 
     result = repo._get_reports_directory()
 
-    assert result.endswith('model_manager/reports')
+    assert result.endswith(os.path.join('model_manager', 'reports'))
     assert os.path.isabs(result)