[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

• LightningDataModule.load_from_checkpoint does not restore subclass from datamodule_hyper_parameters Lightning-AI/pytorch-lightning#21477
• save_hyperparameters(ignore=...) is not persistent across inheritance; ignored params reappear when base class also calls save_hyperparameters Lightning-AI/pytorch-lightning#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
    )

[Copilot]

Pull request overview

This PR introduces a new generalized prediction pipeline intended to work with LightningCLI-saved models/checkpoints, including persisting classification label names into checkpoints for consistent prediction output formatting.

Changes:

• Add checkpoint persistence of classification_labels (derived from a dataset classes.txt) and wire the dataset path into model init via LightningCLI argument linking.
• Introduce a new SMILES prediction entrypoint (chebai/result/prediction.py) that reconstructs model/datamodule from checkpoint hyperparameters.
• Refactor XYBaseDataModule.predict_dataloader to build a prediction dataloader from an in-memory SMILES list, plus update docs/tests and add VS Code workspace files.

Reviewed changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
tests/unit/cli/testCLI.py	Adjusts CLI unit test model args (smaller hidden layer).
tests/unit/cli/mock_dm.py	Adds classes_txt_file_path for CLI linking in tests.
tests/unit/cli/classification_labels.txt	Adds sample classification labels used by CLI unit tests.
chebai/trainer/CustomTrainer.py	Removes prior bespoke prediction logic and overrides predict().
chebai/result/prediction.py	Adds new prediction script/class for SMILES/file inference from checkpoint.
chebai/preprocessing/datasets/base.py	Refactors prediction dataloader flow and adds classes_txt_file_path.
chebai/models/base.py	Adds label-file loading + saving classification_labels into checkpoints.
chebai/cli.py	Links data.classes_txt_file_path into model.init_args.classes_txt_file_path.
README.md	Updates prediction instructions to use the new prediction script.
.vscode/settings.json	Adds VS Code project settings (currently invalid JSON).
.vscode/extensions.json	Adds recommended VS Code extensions.
.gitignore	Stops ignoring the entire .vscode directory (only ignores launch.json).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The new prediction preprocessing path (predict_dataloader returning (DataLoader, valid_indices) plus _process_input_for_prediction / _preprocess_smiles_for_pred) isn't covered by unit tests. Since XYBaseDataModule already has unit tests, adding coverage for valid-index scattering and the empty/all-filtered case would help prevent regressions.

[Copilot]

XYBaseDataModule.predict_dataloader now requires smiles_list (and returns a tuple), which is incompatible with PyTorch Lightning's expected predict_dataloader(self) signature. Any use of Trainer.predict(...) / LightningCLI predict will fail with a missing-argument error. Keep the Lightning-compatible predict_dataloader() signature and add a separate helper (e.g., predict_dataloader_from_smiles(...)) or make smiles_list optional and fall back to the previous behavior when it is None.

[Copilot]

The predictor instantiates a fresh model via instantiate_module(...) but never loads the checkpoint weights (ckpt_file["state_dict"]) into it. As written, inference will run with randomly initialized weights. Load the state_dict into self._model (and consider strict=False with a clear warning if keys mismatch), or use Lightning's load_from_checkpoint(...) so weights are restored correctly.

[Copilot]

self._dm_hparams.pop("splits_file_path") will raise KeyError for checkpoints/datamodules that don't have this key. Use a default (pop(..., None)) or gate this removal on key presence so inference works across datamodules.

Suggested change

	self._dm_hparams.pop("splits_file_path")
	self._dm_hparams.pop("splits_file_path", None)

[Copilot]

CustomTrainer.predict() now unconditionally raises NotImplementedError, which breaks Lightning's predict workflow and the existing chebai CLI predict subcommand (it will always crash). If prediction is intentionally unsupported through the Trainer, remove/disable the CLI predict subcommand; otherwise implement predict() by delegating to super().predict(...) (or a supported prediction path). Also, the error message points to Prediction.* but the new class is named Predictor.

Suggested change

	raise NotImplementedError(
	"CustomTrainer.predict is not implemented."
	"Use `Prediction.predict_from_file` or `Prediction.predict_smiles` from `chebai/result/prediction.py` instead."
	"""
	Runs prediction using the underlying Lightning Trainer implementation.
	This restores the default `Trainer.predict` behavior so that Lightning's
	prediction workflow and the `chebai` CLI `predict` subcommand work as expected.
	"""
	return super().predict(
	model=model,
	dataloaders=dataloaders,
	datamodule=datamodule,
	return_predictions=return_predictions,
	ckpt_path=ckpt_path,

[Copilot]

_process_input_for_prediction treats model_hparams as optional, but then immediately indexes model_hparams["out_dim"]. If model_hparams is omitted/None this will crash. Either make model_hparams required in the public API (and type it accordingly) or add a clear validation/error early in predict_dataloader/_process_input_for_prediction.

[Copilot]

ckpt_file.get("classification_labels") can return None; the next line calls len(self._classification_labels) which will raise TypeError. Handle the missing key explicitly (e.g., default to [] and raise a clear error if absent) so users get an actionable message when using older checkpoints.