train_binary.py checkpoints incompatible with inference_PPI_singleGPU.py

[madnessfish]

Description
Checkpoints saved by train_binary.py cannot be loaded by inference_PPI_singleGPU.py due to architecture mismatch.
Problem

train_binary.py uses AutoModelForSequenceClassification:
pythonself.model = AutoModelForSequenceClassification.from_pretrained(model_name, config=self.config, **automodel_args)
But inference_PPI_singleGPU.py uses PLMinteract which wraps AutoModelForMaskedLM:
pythonself.esm_mask = AutoModelForMaskedLM.from_pretrained(model_name, config=config)

This causes a RuntimeError when loading checkpoints:

RuntimeError: Error(s) in loading state_dict for PLMinteract:
    Missing key(s) in state_dict: "esm_mask.esm.embeddings.word_embeddings.weight", ...
    Unexpected key(s) in state_dict: "esm.embeddings.word_embeddings.weight", ...

Question
What is the intended inference script for checkpoints trained with train_binary.py? Is there a separate inference script I should be using, or should I modify inference_PPI_singleGPU.py to handle both checkpoint formats?

[liudan111]

Thank you for your question. inference_PPI_singleGPU.py is intended for checkpoints trained with train_mlm.py. For models trained with train_binary.py, you can modify inference_PPI_singleGPU.py to load the model using AutoModelForSequenceClassification.

model = AutoModelForSequenceClassification.from_pretrained(
    model_name, config=self.config, **automodel_args
)

load_checkpoint = torch.load(checkpoint_path, map_location="cpu")
model.load_state_dict(load_checkpoint["model"])

model.eval()
for _, (features, labels) in enumerate(dev_dataloader):
    with torch.no_grad():
        outputs = model(**features, return_dict=True)
        logits = activation_fct(outputs.logits)

I will upload an inference script that supports both binary-trained and MLM-trained models.