From 5fa19f4294f5f18a7373ec506a478eca8355c50c Mon Sep 17 00:00:00 2001 From: ncguilbeault Date: Fri, 1 Nov 2024 12:43:11 +0000 Subject: [PATCH] Add custom unpickler to fix issue with missing module --- src/Bonsai.ML.HiddenMarkovModels/main.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/Bonsai.ML.HiddenMarkovModels/main.py b/src/Bonsai.ML.HiddenMarkovModels/main.py index 5cb9af88..9e4a76c3 100644 --- a/src/Bonsai.ML.HiddenMarkovModels/main.py +++ b/src/Bonsai.ML.HiddenMarkovModels/main.py @@ -11,6 +11,11 @@ npr.seed(0) +class CustomUnpickler(pickle.Unpickler): + def find_class(self, module, name): + if module == 'main' and name == 'HiddenMarkovModel': + return HiddenMarkovModel + return super().find_class(module, name) class HiddenMarkovModel(HMM): @@ -158,7 +163,7 @@ def save_model(self, path: str): @classmethod def load_model(cls, path: str): with open(path, 'rb') as f: - return pickle.load(f) + return CustomUnpickler(f).load() def fit_async(self, observation: list[float],