diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 5478103..b3473c6 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -22,250 +22,6 @@ class MLFlowRepository(): self.model_serving = ModelServing(tracking_uri=host, username=username, password=password) -<<<<<<< HEAD -======= - def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str): - """ - Get the current data as a DataFrame and update the prediction and target columns - - Parameters: - current_data (pd.DataFrame): the current data - model_name (str): the name of the model - target (str): the target column - - Returns: - DataFrame: the current data as a DataFrame - - - """ - predictions = current_data['prediction'] - - target = current_data[target] - current_data = self.model_serving.get_transformed_data( - model_name, current_data, by='model') - current_data['prediction'] = predictions - current_data['target'] = target - - return pd.DataFrame(current_data).dropna() - - def get_artifact(self, destination: str, search_by: str, run_id: str = None, - model_name: str = None, artifact_name: str = None) -> None: - """ - Get an artifact in MLflow by experiment or model and save it to a - destination path using API. - If the artifact is searched by model, the latest production version will be used. - - Args: - destination: The destination path to save the artifact. - search_by: The way to search for the artifact ('experiment' or 'model'). - run_id: The run ID of the experiment (if search_by is "experiment"). - model_name: The name of the model (if search_by is "model"). - artifact_name: The path of the artifact to download. - - Returns: - artifact: The artifact(.csv) downloaded from MLflow. - """ - - self.model_serving.get_artifact(destination=destination, search_by=search_by, - run_id=run_id, model_name=model_name, artifact_name=artifact_name) - - def calculate_model_metrics(self, real_data, predictions, flag): - """ - Function to calculate the metrics of a model using API - - Parameters: - real_data (array): the real data - predictions (array): the predictions - - Returns: - dict: the metrics of the model including MSE and R2 - """ - return self.model_serving.get_model_metrics( - reference_data=None, real_data=real_data, - predictions=predictions, type_flag=flag) - - def get_experiment_by_run_id(self, run_id: str) -> dict: - # Get the run information using the run_id - run = mlflow.get_run(run_id) - - # Extract the experiment ID from the run - experiment_id = run.info.experiment_id - - # Get the experiment details using the experiment ID - experiment = mlflow.get_experiment(experiment_id) - experiment_name = experiment.name - return experiment_name - - def get_next_run_name(self, model_name: str) -> str: - """ - Function to get the next run number of a specific model - - Parameters: - model_name (str): the name of the model - - Returns: - str: the next run number - """ - - runs = mlflow.search_runs( - experiment_names=[model_name], order_by=["start_time desc"]) - next_run_number = len(runs) + 1 - return f"{model_name}-{next_run_number}" - - def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: - """ - Retrain a model with new data. - - Parameters: - - data (pandas.DataFrame): The new data to use for retraining. - - model_name (str): The name of the model to retrain. - - metrics_list (list): The metrics to be used to compare the models. - - compare_metrics (bool): If True, the retrain will only be considered if - the new model is better than the current one. - If False, the retrain will always be considered. - - split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets. - If False, the data will be used as a unique block for retraining. - - update_report (bool): If True, a report will be created with the data of the retrained model. - - update_transformation (bool): If True, the model will be updated in the MLflow tracking server. - - update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server. - - shuffle_data (bool): If True, the data will be shuffled before splitting. - - model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'. - - - Returns: - - mlflow.sklearn.Model: The retrained prediction model. - - mlflow.sklearn.Model: The retrained data model. - - mse (float): The mean squared error of the retrained model. - - r2 (float): The R-squared score of the retrained model. - """ - - # load predictor model - predictor_uri = f"models:/{model_name}/production" - # load transform model - latest_production_id = self.model_serving.get_model_run_id( - model_name, stage="Production" - ) - transform_uri = self.model_serving.get_model_uri( - latest_production_id, prediction=False - ) - # load - data_model = mlflow.sklearn.load_model(transform_uri) - prediction_model = mlflow.sklearn.load_model(predictor_uri) - data_model = data_model.fit(data) - treated_data = data_model.predict(data) - # align target column with treated_data - target_name = data_model.target_variable - y = data[target_name] - treated_data = pd.merge( - treated_data, y, left_index=True, right_index=True) - prediction_model = prediction_model.fit(treated_data) - # Example usage - experiment = self.get_experiment_by_run_id(latest_production_id) - pred_model_atributes = vars(prediction_model) # load class attributes - data_model_atributes = vars(data_model) # load class attributes - mlflow.set_experiment(experiment) - experiment_description = "Retrain model {model_name} with new data" - current_run_name = self.get_next_run_name(experiment) - with mlflow.start_run( - run_name=current_run_name, description=experiment_description - ) as _run: - # update transfomation model - # fixed parameters - for name_atribute, val_atribute in pred_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # update prediction model - for name_atribute, val_atribute in data_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # dynamic parameters, including model itself - mlflow.sklearn.log_model(data_model, "data_model") - file_path = f"laborious/data/raw_data_{model_name}.csv" - data.to_csv( - f"laborious/data/raw_data_{model_name}.csv", index=True) - # log the data raw - mlflow.log_artifact(file_path) - - # dynamic parameters, including model itself - mlflow.sklearn.log_model(prediction_model, "prediction_model") - mlflow.log_param("retrain", True) - - return "Model retrained successfully", experiment - - def get_experiment(self, experiment_name: str) -> int: - experiment = mlflow.get_experiment_by_name(experiment_name) - - if experiment is None: - raise ValueError(f'Experiment {experiment_name} not found') - - return int(experiment.experiment_id) - - def get_experiment_last_run(self, experiment_id: int) -> str: - runs = mlflow.search_runs( - experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" - ) - - # Filtrar apenas as runs onde params.retrain == True - filtered_runs = runs[runs["params.retrain"] == 'True'] - - # Converter a coluna 'end_time' para datetime - filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) - - # Ordenar o DataFrame de forma descendente pela coluna 'end_time' - filtered_runs = filtered_runs.sort_values( - by='end_time', ascending=False) - - # Pegar a última run_id do DataFrame filtrado e ordenado - latest_run_id = filtered_runs.iloc[0]['run_id'] - - return latest_run_id - - def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict: - # Registrar o modelo - # Aqui estamos assumindo que você já tem um modelo salvo, - # caso contrário você precisará treiná-lo e salvá-lo primeiro. - # Se o modelo já está registrado, você pode usar o método - # register_model() ou pyfunc.load_model() para isso. - mlflow.register_model( - f"runs:/{run_id}/prediction_model", model_name) - - # Colocar a versão do modelo em produção - # Depois de registrar o modelo, precisamos pegar a versão mais - # recente do modelo e movê-lo para o estágio 'Production' - client = mlflow.tracking.MlflowClient() - - # Obter a versão mais recente registrada do modelo - model_versions = client.get_registered_model( - model_name).latest_versions - max_version = max(model_versions, key=lambda x: int(x.version)).version - - # Mover a versão mais recente do modelo para o estágio de 'Production' - client.transition_model_version_stage( - name=model_name, - version=max_version, - stage="Production", - archive_existing_versions=True - ) - - return { - 'model_name': model_name, - 'version': max_version, - 'mlflow_run_id': run_id - } - - def update_production_model(self, experiment: str, model_name: str) -> dict: - - experiment_id = self.get_experiment(experiment) - run_id = self.get_experiment_last_run(experiment_id) - metadata = self.update_production_model_by_run_id(run_id, model_name) - - metadata['mlflow_experiment_id'] = experiment_id - - return metadata - ->>>>>>> main def transform(self, model_name: str, data: pd.DataFrame, model_retention: int): """ Transform data using a model.