build bert: build does not load model

[Alireza3242]

System Info

a100

Who can help?

@byshiue

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

My bert model is model.safetensors. But examples/bert/build.py cant read it.
I changed this part of code and worked:

if args.model_dir is not None and os.path.exists(
                os.path.join(args.model_dir, "pytorch_model.bin")):
            state_dict = torch.load(
                os.path.join(args.model_dir, "pytorch_model.bin"))
            hf_bert.load_state_dict(state_dict, strict=False)

to

if args.model_dir is not None and os.path.exists(
                os.path.join(args.model_dir, "pytorch_model.bin")):
            state_dict = torch.load(
                os.path.join(args.model_dir, "pytorch_model.bin"))
            hf_bert.load_state_dict(state_dict, strict=False)
elif args.model_dir is not None and os.path.exists(
                os.path.join(args.model_dir, "model.safetensors")):
            state_dict = safetensors.torch.load_file(os.path.join(args.model_dir, "model.safetensors"))
            hf_bert.load_state_dict(state_dict, strict=False)

Expected behavior

read model

actual behavior

not read model

additional notes

nothing

[symphonylyh]

fixed together with your other issue #2373. Thank you!!
I used a try-catch instead:

from safetensors.torch import load_file
...
if args.model_dir is not None:
            try:
                state_dict = torch.load(
                    os.path.join(args.model_dir, "pytorch_model.bin"))
            except FileNotFoundError:
                state_dict = load_file(os.path.join(args.model_dir, "model.safetensors"))
            hf_bert.load_state_dict(state_dict, strict=False)

[tkhanipov]

@symphonylyh
Note: the change only covers the case elif args.model == 'BertForSequenceClassification' or args.model == 'RobertaForSequenceClassification'. The if args.model == 'BertModel' or args.model == 'RobertaModel' case remains buggy. Please have another look at #2187, there is a fix for the latter case (the corresponding bug: #2197).

[tkhanipov]

@symphonylyh Here you said that you were planning to merge the PR