[BUG] `load_model_trainer_states_from_checkpoint()` function is broken

[Legoclones]

Bug description

The transformers4rec.torch.trainer.load_model_trainer_states_from_checkpoint() function can only load t4rec_model_class.pkl, all the other loads are broken. The utils.serialization.load() function is called 4 times in the function, but only one of them successfully loads, the other throw errors. I believe this is because load() is expecting an _io.BufferedReader object (resulting from open()) instead of a path.

Correctly implemented:

Transformers4Rec/transformers4rec/torch/trainer.py

Lines 744 to 746 in 5d59d14

	model = load(
	open(os.path.join(checkpoint_path, "t4rec_model_class.pkl"), "rb")
	)

Incorrectly implemented:

Transformers4Rec/transformers4rec/torch/trainer.py

Line 753 in 5d59d14

load(os.path.join(checkpoint_path, "pytorch_model.bin"), torch_load=True)

Transformers4Rec/transformers4rec/torch/trainer.py

Line 757 in 5d59d14

checkpoint_rng_state = load(rng_file, torch_load=True)

Transformers4Rec/transformers4rec/torch/trainer.py

Line 766 in 5d59d14

load(os.path.join(checkpoint_path, "scaler.pt"), torch_load=True)

Steps/Code to reproduce bug

Run the following code with a Model checkpoint at /checkpoint:

from transformers4rec import torch as tr

trainer = tr.trainer.Trainer(model=tr.model.base.Model(), args=None)
trainer.load_model_trainer_states_from_checkpoint('/checkpoint')

Expected behavior

The function call shouldn't be erroring.

Environment details

• Transformers4Rec version: 23.12.0
• Platform: Linux
• Python version: 3.10.12
• Huggingface Transformers version: 4.27.1
• PyTorch version (GPU?): 2.1.2 (no)

[shaun0927]

Quick follow-up after re-checking main (HEAD @ 8bf122f5, post-#807/#808/#809).

The commits that shipped after this issue was filed do not restore working file-I/O checkpoint loading — they change the shape of the breakage. Specifically:

1. Fix scaler #807 (0e31f575) rewrote load_model_trainer_states_from_checkpoint away from path-based loading. The new signature is (self, model, scaler=None) — no path argument at all. The body of the function no longer reads pytorch_model.bin, rng_state.pth, or scaler.pt from disk. Callers that still pass a checkpoint path (the API shape this very issue reports) now silently bind the string to model and the method returns without loading anything. There is no deprecation warning or type check.

2. Fix scaler #808 (41b14d7b) + Fix scaler #809 (8bf122f5) rewrote Model.load, but Fix scaler #809 introduced a regression that makes the new API unreachable:

   Transformers4Rec/transformers4rec/torch/model/base.py

   Line 917 in 8bf122f

   if isinstance(state_dict, Dict[str, torch.Tensor]):

   if isinstance(state_dict, Dict[str, torch.Tensor]):
       model.load_state_dict(state_dict, strict=strict)
   else:
       raise ValueError(\"\`state_dict\` must be a dictionary of parameter (torch) tensors.\")

   On Python 3.9+, isinstance(x, Dict[str, torch.Tensor]) raises TypeError: Subscripted generics cannot be used with class and instance checks — so every Model.load(...) call errors out before the dispatch, regardless of how the state dict was obtained. I filed a separate issue for this ([BUG] Model.load raises TypeError on Python 3.9+ (parameterized generic in isinstance) #810).

3. utils/serialization.load(..., torch_load=True) is now a dead kwarg. PR Fix scaler #807 removed the if torch_load: result = torch.load(file) branch but kept torch_load=False in the function signature. External callers that still pass torch_load=True receive restricted-Unpickler output (wrong type) with no warning. PR Fix scaler #808 then deleted the module entirely, leaving any caller of from transformers4rec.utils.serialization import load with an ImportError.

So the situation from the user's perspective (reproduced against main):

• Original issue (path-based load_model_trainer_states_from_checkpoint): the function no longer loads anything from a path. Call sites that worked on v23.12.00 silently do nothing.
• New-style load (Model.load(state_dict, heads)): immediately TypeError due to [BUG] Model.load raises TypeError on Python 3.9+ (parameterized generic in isinstance) #810.

A minimal restoration that does not re-introduce the pickle-RCE surface would look like:

def load_model_trainer_states_from_checkpoint(self, checkpoint_path, model):
    import os, random, numpy as np, torch
    self.model = model
    self.model.load_state_dict(
        torch.load(os.path.join(checkpoint_path, \"pytorch_model.bin\"),
                   weights_only=True)
    )
    rng = torch.load(os.path.join(checkpoint_path, \"rng_state.pth\"),
                     weights_only=True)
    random.setstate(rng[\"python\"])
    np.random.set_state(rng[\"numpy\"])
    torch.random.set_rng_state(rng[\"cpu\"])
    torch.cuda.random.set_rng_state_all(rng[\"cuda\"])
    if self._use_cuda_amp:
        self.scaler.load_state_dict(
            torch.load(os.path.join(checkpoint_path, \"scaler.pt\"),
                       weights_only=True)
        )

Using weights_only=True avoids the pickle-RCE concern that motivated PR #802 in the first place. I can send a PR once #810 is addressed (fixing #810 is a prerequisite since Model.load is on the common load path).