diff --git a/dataikuapi/dss/ml.py b/dataikuapi/dss/ml.py index dac73af4..98d72e7a 100644 --- a/dataikuapi/dss/ml.py +++ b/dataikuapi/dss/ml.py @@ -1752,6 +1752,84 @@ def get_diagnostics(self): diagnostics = self.details.get("trainDiagnostics", {}) return [DSSMLDiagnostic(d) for d in diagnostics.get("diagnostics", [])] + def generate_documentation(self, folder_id=None, path=None): + """ + Start the model document generation from a template docx file in a managed folder, + or from the default template if no folder id and path are specified. + + :param folder_id: (optional) the id of the managed folder + :param path: (optional) the path to the file from the root of the folder + :return: A :class:`~dataikuapi.dss.future.DSSFuture` representing the model document generation process + """ + if bool(folder_id) != bool(path): + raise ValueError("Both folder id and path arguments are required to use a template from folder. Use without argument to generate the model documentation using the default template") + + template_mode_url = "default-template" if folder_id is None and path is None else "template-in-folder" + + if self.mltask is not None: + f = self.mltask.client._perform_json( + "POST", "/projects/%s/models/lab/%s/%s/models/%s/generate-documentation-from-%s" % + (self.mltask.project_key, self.mltask.analysis_id, self.mltask.mltask_id, self.mltask_model_id, template_mode_url), + params={"folderId": folder_id, "path": path}) + return DSSFuture(self.mltask.client, f["jobId"]) + else: + f = self.saved_model.client._perform_json( + "POST", "/projects/%s/savedmodels/%s/versions/%s/generate-documentation-from-%s" % + (self.saved_model.project_key, self.saved_model.sm_id, self.saved_model_version, template_mode_url), + params={"folderId": folder_id, "path": path}) + return DSSFuture(self.saved_model.client, job_id=f["jobId"]) + + def generate_documentation_from_custom_template(self, fp): + """ + Start the model document generation from a docx template (as a file object). + + :param object fp: A file-like object pointing to a template docx file + :return: A :class:`~dataikuapi.dss.future.DSSFuture` representing the model document generation process + """ + files = {'file': fp} + if self.mltask is not None: + f = self.mltask.client._perform_json( + "POST", "/projects/%s/models/lab/%s/%s/models/%s/generate-documentation-from-custom-template" % + (self.mltask.project_key, self.mltask.analysis_id, self.mltask.mltask_id, self.mltask_model_id), + files=files) + return DSSFuture(self.mltask.client, f["jobId"]) + else: + f = self.saved_model.client._perform_json( + "POST", "/projects/%s/savedmodels/%s/versions/%s/generate-documentation-from-custom-template" % + (self.saved_model.project_key, self.saved_model.sm_id, self.saved_model_version), + files=files) + return DSSFuture(self.saved_model.client, job_id=f["jobId"]) + + def download_documentation_stream(self, export_id): + """ + Download a model documentation, as a binary stream. + + Warning: this stream will monopolize the DSSClient until closed. + + :param export_id: the id of the generated model documentation returned as the result of the future + :return: A :class:`~dataikuapi.dss.future.DSSFuture` representing the model document generation process + """ + if self.mltask is not None: + return self.mltask.client._perform_raw( + "GET", "/projects/%s/models/lab/documentations/%s" % (self.mltask.project_key, export_id)) + else: + return self.saved_model.client._perform_raw( + "GET", "/projects/%s/savedmodels/documentations/%s" % (self.saved_model.project_key, export_id)) + + def download_documentation_to_file(self, export_id, path): + """ + Download a model documentation into the given output file. + + :param export_id: the id of the generated model documentation returned as the result of the future + :param path: the path where to download the model documentation + :return: None + """ + stream = self.download_documentation_stream(export_id) + with open(path, 'wb') as f: + for chunk in stream.iter_content(chunk_size=10000): + if chunk: + f.write(chunk) + f.flush() class DSSMLDiagnostic(object): """ @@ -2462,6 +2540,7 @@ def get_scoring_pmml_stream(self): "GET", "/projects/%s/savedmodels/%s/versions/%s/scoring-pmml" % (self.saved_model.project_key, self.saved_model.sm_id, self.saved_model_version)) + ## Post-train computations def compute_subpopulation_analyses(self, split_by, wait=True, sample_size=1000, random_state=1337, n_jobs=1, debug_mode=False):