[Feature]: Generalize Prediction pipeline for Lightning CLI models

[aditya0by0]

Generalize prediction logic

Please merge below PRs after this PR:

• Prediction functional for Graphs python-chebai-graph#31
• Utilize new prediction pipeline python-chebifier#25
• New prediction pipeline chebifier-web#28

Related Discussion

• Reliability of `_class_path` in Lightning checkpoints when using `instantiate_module` Lightning-AI/pytorch-lightning#21474

Related bugs rectified in Lightning for the pipeline

• https://github.com/Lightning-AI/pytorch-lightning/ issues/21477
• https://github.com/Lightning-AI/pytorch-lightning/ issues/21488

Additional changes

• Save class labels in checkpoint under the key "classification_labels"

  >>> dt = torch.load("path/to/file.ckpt", weights_only=False, map_location="cpu")
  >>> dt.keys()
  dict_keys(['epoch', 'global_step', 'pytorch-lightning_version', 'state_dict', 'loops', 'callbacks', 'optimizer_states', 'lr_schedulers', 'hparams_name', 'hyper_parameters', 'datamodule_hparams_name', 'datamodule_hyper_parameters', 'classification_labels'])

• Wrap inference with torch.inference_mode() to avoid gradient tracking (see Avoid gradient tracking python-chebifier#21)

  • Behavior of model.eval() in PyTorch
  • Difference between torch.no_grad() and torch.inference_mode()

• Use torch.compile for faster inference

  • https://huggingface.co/docs/transformers/en/perf_torch_compile
  • https://docs.pytorch.org/tutorials/intermediate/torch_compile_full_example.html
  • https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

[aditya0by0]

@sfluegel05,

Could you confirm our agreed approach for handling old_checkpoint files which don't classification labels stored in them?

1. Update code to handle legacy checkpoints: This requires adding logic to prediction.py and the chebifier repo to ingest external class files.

   • Concerns: Adds boilerplate and permanent complexity to handle a temporary issue.

2. Patch old checkpoints (Preferred): Use the below one-time script to inject labels into the existing files.

   • Benefits: Keeps the codebase clean and ensures all checkpoints follow a standardized schema.

I’m in favor of Option 2 to avoid carrying technical debt in the prediction logic. Does this match your understanding?

I'm willing to add this script to the repo, and small readme note for old checkpoints for option 2.

import sys

import torch


def add_class_labels_to_checkpoint(input_path, classes_file_path):
    with open(classes_file_path, "r") as f:
        class_labels = [line.strip() for line in f.readlines()]

    assert len(class_labels) > 0, "The classes file is empty."

    # 1. Load the checkpoint
    checkpoint = torch.load(
        input_path, map_location=torch.device("cpu"), weights_only=False
    )

    if "classification_labels" in checkpoint:
        print(
            "Warning: 'classification_labels' key already exists in the checkpoint and will be overwritten."
        )


    # 2. Add your custom key/value pair
    checkpoint["classification_labels"] = class_labels

    # 3. Save the modified checkpoint
    output_path = input_path.replace(".ckpt", "_modified.ckpt")
    torch.save(checkpoint, output_path)
    print(f"Successfully added classification_labels and saved to {output_path}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python modify_checkpoints.py <input_checkpoint> <classes_file>")
        sys.exit(1)

    input_ckpt = sys.argv[1]
    classes_file = sys.argv[2]

    add_class_labels_to_checkpoint(
        input_path=input_ckpt, classes_file_path=classes_file
    )