mirror of
https://github.com/luxfhe/ml-sdk.git
synced 2026-10-30 18:13:04 +00:00
zeekay
bde2e2531b
Machine Learning with FHE - fork of Zama Concrete-ML - Built-in FHE-friendly ML models - scikit-learn compatible API - PyTorch model conversion - Quantization-aware training
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
"""Finalize Jupyter notebooks."""
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
# pylint: disable=too-many-nested-blocks
|
|
def main():
|
|
"""Finalize
|
|
|
|
Raises:
|
|
Exception: if notebooks are not sanitized.
|
|
"""
|
|
|
|
parser = argparse.ArgumentParser(description="Sanitizer for Jupyter Notebooks")
|
|
|
|
parser.add_argument("base", type=str, nargs="+", help="directory which contains the notebooks")
|
|
parser.add_argument("--check", action="store_true", help="flag to enable just checking mode")
|
|
|
|
args = parser.parse_args()
|
|
|
|
for directory in args.base:
|
|
base = Path(directory)
|
|
notebooks = base.glob("**/*.ipynb")
|
|
|
|
for notebook in notebooks:
|
|
path = str(notebook)
|
|
|
|
if "_build" in path or ".ipynb_checkpoints" in path:
|
|
continue
|
|
|
|
with open(notebook, "r", encoding="utf-8") as f:
|
|
try:
|
|
content = json.load(f)
|
|
except Exception as exception:
|
|
raise Exception(f"Error while processing {notebook=}") from exception
|
|
|
|
if args.check:
|
|
try:
|
|
metadata = content["metadata"]
|
|
assert len(metadata) == 1
|
|
assert "execution" in metadata
|
|
|
|
execution = metadata["execution"]
|
|
assert len(execution) == 1
|
|
assert "timeout" in execution
|
|
|
|
timeout = execution["timeout"]
|
|
assert timeout == 10800 # 3 hours
|
|
|
|
for cell in content["cells"]:
|
|
if "metadata" in cell:
|
|
assert len(cell["metadata"]) == 0
|
|
|
|
except Exception as exception:
|
|
message = "Notebooks are not sanitized. Please run `make conformance`."
|
|
print(message)
|
|
raise Exception(message) from exception
|
|
else:
|
|
content["metadata"] = {
|
|
"execution": {
|
|
"timeout": 10800, # 3 hours
|
|
}
|
|
}
|
|
|
|
for cell in content["cells"]:
|
|
if "metadata" in cell:
|
|
cell["metadata"] = {}
|
|
|
|
with open(notebook, "w", newline="\n", encoding="utf-8") as f:
|
|
json.dump(content, f, indent=1, ensure_ascii=False)
|
|
f.write("\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|