unsloth/docker/unsloth_nb_content_sig.py
2026-06-16 03:49:07 +00:00

119 lines
3.9 KiB
Python

#!/usr/bin/env python3
# Compare the *content* of two Unsloth notebooks, ignoring the auto-generated
# top/bottom boilerplate that update_all_notebooks.py stamps on every notebook.
#
# Every generated notebook has the same shape:
# - a "top" of boilerplate: the "To run this, press Runtime" announcement, the
# "### News" / Unsloth Studio announcement cells, and the %%capture install
# cell. These churn constantly (new pip pins, new announcements, new links).
# - the "middle": the actual tutorial (data prep, train, inference, save).
# - a "bottom": the "And we're done ... licensed LGPL-3.0" footer cell.
#
# The boot-time notebook refresh uses this to avoid rewriting a user's notebook
# when only that boilerplate moved upstream. We hash ONLY the middle (the cells
# that are not install/announcement/footer) and compare. Outputs, execution
# counts, cell ids and metadata are ignored, so merely running a notebook never
# changes the signature.
#
# Usage:
# unsloth_nb_content_sig.py <a.ipynb> <b.ipynb> -> prints SAME | DIFF | ERR
# unsloth_nb_content_sig.py <a.ipynb> -> prints the middle digest
#
# Exit code is always 0; the decision is the printed word. On any parse problem
# we print ERR / nothing so the caller can fall back to its whole-file logic.
import hashlib
import json
import sys
# Lowercased substrings that mark a markdown cell as top/bottom boilerplate.
_BOILERPLATE_MD = (
"to run this, press", # Colab/AMD run announcement
'press "*runtime*"',
"### news", # News heading
"introducing **unsloth studio**", # rotating announcement body
"you will learn how to do", # announcement tail
"this notebook is licensed", # announcement license line
"and we're done", # footer opener
"this notebook and all unsloth notebooks are licensed", # footer license
"join discord if you need help", # footer
"star us on", # footer
"some other resources", # footer resources block
)
def _text(cell):
src = cell.get("source", "")
if isinstance(src, list):
src = "".join(src)
return src.replace("\r\n", "\n").replace("\r", "\n")
def _is_install_code(cell):
if cell.get("cell_type") != "code":
return False
t = _text(cell)
low = t.lower()
if "pip install" in low or "pip3-autoremove" in low:
return True
first = t.lstrip().split("\n", 1)[0].strip().lower()
return first.startswith("%%capture") or first.startswith("%%bash")
def _is_boilerplate_md(cell):
if cell.get("cell_type") != "markdown":
return False
low = _text(cell).lower()
return any(m in low for m in _BOILERPLATE_MD)
def _is_boilerplate(cell):
return _is_install_code(cell) or _is_boilerplate_md(cell)
def middle_digest(path):
"""sha256 over the (type, source) of every non-boilerplate cell, or None."""
try:
with open(path, "r", encoding = "utf-8") as f:
nb = json.load(f)
except Exception:
return None
cells = nb.get("cells")
if not isinstance(cells, list):
return None
h = hashlib.sha256()
for cell in cells:
if not isinstance(cell, dict):
continue
if _is_boilerplate(cell):
continue
h.update(b"\x00")
h.update(str(cell.get("cell_type", "")).encode("utf-8"))
h.update(b"\x01")
h.update(_text(cell).encode("utf-8"))
return h.hexdigest()
def main(argv):
if len(argv) == 2:
d = middle_digest(argv[1])
if d is None:
print("ERR")
return 0
print(d)
return 0
if len(argv) == 3:
a = middle_digest(argv[1])
b = middle_digest(argv[2])
if a is None or b is None:
print("ERR")
elif a == b:
print("SAME")
else:
print("DIFF")
return 0
print("ERR")
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv))