[Bug] Unable to save model in HDF5 format during learning rate search

[nhansendev]

When trying to use learning rate search with the standard preprocessor and NBSVM model for text classification I encountered an error related to saving the model in HDF5 format (done automatically in lr_find). It warns that the HDF5 format is legacy (currently using tensorflow 2.15.1), but fails to load the file afterwards suggesting that it failed to actually save as well.

Strangely, there were no issues using this same setup with 'distilbert' as the pre-processor instead.

Example code:

self.train_data, self.val_data, self.preprocess = ktrain.text.texts_from_df(
  training_data,
  "combined",
  label_columns=["label"],
  val_df=None,
  max_features=20000,
  maxlen=512,
  val_pct=0.1,
  ngram_range=1,
  preprocess_mode='standard',
  random_state=random_seed,
  verbose=0
)

self.model = ktrain.text.text_classifier('nbsvm', self.train_data, preproc=self.preprocess)
self.learner = ktrain.get_learner(self.model, train_data=self.train_data, val_data=self.val_data, batch_size=16)
self.learner.lr_find(max_epochs=6, show_plot=True)

Output:

Is Multi-Label? False
Is Multi-Label? False
compiling word ID features...
maxlen is 512
building document-term matrix... this may take a few moments...
rows: 1-1800
computing log-count ratios...
done.
simulating training for different learning rates... this may take a few moments...
/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/tf_keras/src/engine/training.py:3098: UserWarning: You are saving your model as an HDF5 file via `model.save()`. This file format is considered legacy. We recommend using instead the native TF-Keras format, e.g. `model.save('my_model.keras')`.
  saving_api.save_model(
Epoch 1/6
113/113 [==============================] - 9s 9ms/step - loss: 0.6944 - accuracy: 0.4967
Epoch 2/6
113/113 [==============================] - 1s 8ms/step - loss: 0.6932 - accuracy: 0.5006
Epoch 3/6
113/113 [==============================] - 1s 8ms/step - loss: 0.6698 - accuracy: 0.5956
Epoch 4/6
113/113 [==============================] - 1s 8ms/step - loss: 0.4787 - accuracy: 0.8456
Epoch 5/6
113/113 [==============================] - 1s 8ms/step - loss: 0.2304 - accuracy: 0.9333
Epoch 6/6
113/113 [==============================] - 1s 5ms/step - loss: 2.3402 - accuracy: 0.8648
Call to keras.models.load_model failed. Try manually invoking this function to investigate error and report issue if necessary.
Process Process-6:
Traceback (most recent call last):
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/ktrain/core.py", line 2048, in _load_model
    model = keras.models.load_model(
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/tf_keras/src/saving/saving_api.py", line 262, in load_model
    return legacy_sm_saving_lib.load_model(
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/tf_keras/src/saving/legacy/save.py", line 229, in load_model
    with tf.__internal__.load_context(options):
AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' has no attribute 'load_context'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/nate/DTSA5798/model.py", line 34, in _run
    mod.estimate_lr()
  File "/home/nate/DTSA5798/model.py", line 73, in estimate_lr
    self.learner.lr_find(max_epochs=6, show_plot=True)
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/ktrain/core.py", line 653, in lr_find
    self.load_model(temp_folder)
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/ktrain/core.py", line 394, in load_model
    self.model = _load_model(
  File "/home/nate/DTSA5798/.venv/lib/python3.10/site-packages/ktrain/core.py", line 2055, in _load_model
    raise Exception("Error detected: %s" % (e))
Exception: Error detected: module 'tensorflow._api.v2.compat.v2.__internal__' has no attribute 'load_context'