Module brevettiai.utils.model_version
Expand source code
import os
import json
import tarfile
import hashlib
from io import BytesIO
def get_model_version(obj, blength=4):
"""
Calculate version of model archive
:param obj: fname path to model archive or BufferedIOBase
:param blength: bytelength of model version (4)
:return:
"""
try:
hash_md5 = hashlib.md5()
if isinstance(obj, BytesIO):
for chunk in iter(lambda: obj.read(4096), b""):
hash_md5.update(chunk)
else:
with open(obj, "rb") as f:
for chunk in iter(lambda: f.read(4096), b""):
hash_md5.update(chunk)
return int.from_bytes(hash_md5.digest()[:blength], byteorder='little', signed=False)
except FileNotFoundError:
return -1
def check_model_version(fname=None, version=None, fileobj=None, blength=4):
"""
Check model version is correct
:param fname: path to model archive, or filelike obj
:param fileobj: BytesIO if model in memory
:param blength: bytelength of model version (4)
:return: Version number of file if correct, False if not correct and None if not a known file type
"""
try:
fileobj = fileobj or fname
version = version or int(fname.rsplit(".", 3)[-3])
true_version = get_model_version(fileobj, blength)
return true_version if version == true_version else False
except Exception:
return False
def saved_model_model_meta_filename(saved_model_dir):
return os.path.join(saved_model_dir, "assets.extra", "criterion_model_meta.json")
def package_saved_model(saved_model_dir, name="saved_model", output=None, model_meta=None):
"""
Utility function to package saved model to archive
:param saved_model_dir: path to saved model directory
:param name: name of model
:param output: forced output path use given name in place of name with model version
:return: archive path
"""
if model_meta is not None:
model_meta_path = saved_model_model_meta_filename(saved_model_dir)
os.makedirs(os.path.dirname(model_meta_path), exist_ok=True)
with open(model_meta_path, "w") as fp:
json.dump(model_meta, fp)
# Package
saved_model_tar = output or os.path.join(saved_model_dir + '.tar.gz')
with tarfile.open(saved_model_tar, "w:gz") as tar:
tar.add(saved_model_dir, arcname="saved_model")
# Rename
if output is None:
root_dir = os.path.dirname(saved_model_dir)
output = os.path.join(root_dir, f"{name}.{get_model_version(saved_model_tar)}.tar.gz")
os.rename(saved_model_tar, output)
return output
Functions
def check_model_version(fname=None, version=None, fileobj=None, blength=4)
-
Check model version is correct :param fname: path to model archive, or filelike obj :param fileobj: BytesIO if model in memory :param blength: bytelength of model version (4) :return: Version number of file if correct, False if not correct and None if not a known file type
Expand source code
def check_model_version(fname=None, version=None, fileobj=None, blength=4): """ Check model version is correct :param fname: path to model archive, or filelike obj :param fileobj: BytesIO if model in memory :param blength: bytelength of model version (4) :return: Version number of file if correct, False if not correct and None if not a known file type """ try: fileobj = fileobj or fname version = version or int(fname.rsplit(".", 3)[-3]) true_version = get_model_version(fileobj, blength) return true_version if version == true_version else False except Exception: return False
def get_model_version(obj, blength=4)
-
Calculate version of model archive :param obj: fname path to model archive or BufferedIOBase :param blength: bytelength of model version (4) :return:
Expand source code
def get_model_version(obj, blength=4): """ Calculate version of model archive :param obj: fname path to model archive or BufferedIOBase :param blength: bytelength of model version (4) :return: """ try: hash_md5 = hashlib.md5() if isinstance(obj, BytesIO): for chunk in iter(lambda: obj.read(4096), b""): hash_md5.update(chunk) else: with open(obj, "rb") as f: for chunk in iter(lambda: f.read(4096), b""): hash_md5.update(chunk) return int.from_bytes(hash_md5.digest()[:blength], byteorder='little', signed=False) except FileNotFoundError: return -1
def package_saved_model(saved_model_dir, name='saved_model', output=None, model_meta=None)
-
Utility function to package saved model to archive :param saved_model_dir: path to saved model directory :param name: name of model :param output: forced output path use given name in place of name with model version :return: archive path
Expand source code
def package_saved_model(saved_model_dir, name="saved_model", output=None, model_meta=None): """ Utility function to package saved model to archive :param saved_model_dir: path to saved model directory :param name: name of model :param output: forced output path use given name in place of name with model version :return: archive path """ if model_meta is not None: model_meta_path = saved_model_model_meta_filename(saved_model_dir) os.makedirs(os.path.dirname(model_meta_path), exist_ok=True) with open(model_meta_path, "w") as fp: json.dump(model_meta, fp) # Package saved_model_tar = output or os.path.join(saved_model_dir + '.tar.gz') with tarfile.open(saved_model_tar, "w:gz") as tar: tar.add(saved_model_dir, arcname="saved_model") # Rename if output is None: root_dir = os.path.dirname(saved_model_dir) output = os.path.join(root_dir, f"{name}.{get_model_version(saved_model_tar)}.tar.gz") os.rename(saved_model_tar, output) return output
def saved_model_model_meta_filename(saved_model_dir)
-
Expand source code
def saved_model_model_meta_filename(saved_model_dir): return os.path.join(saved_model_dir, "assets.extra", "criterion_model_meta.json")