feat: WSL meta arch

[xrusnack]

No description provided.

[coderabbitai]

Warning

Rate limit exceeded

@xrusnack has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 36 minutes and 18 seconds before requesting another review.

Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 36 minutes and 18 seconds.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9f85c285-7827-49cc-9a47-0b9dc6b92b96

📥 Commits

Reviewing files that changed from the base of the PR and between fcce74f and fb9f2b8.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (6)

• configs/experiment/modeling/training/nuclei_level/prostate_cancer_mmci_tl.yaml
• configs/experiment/modeling/training/nuclei_level/radboud.yaml
• configs/model/meta_archs/nuclei_wsl.yaml
• nuclei_graph/__init__.py
• nuclei_graph/nuclei_wsl_meta_arch.py
• pyproject.toml

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/meta-arch/wsl

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new weakly supervised learning (WSL) meta-architecture designed for analyzing nuclei graphs. It provides the core implementation for training, validating, and testing this model using PyTorch Lightning, along with specific configuration files to facilitate its application to different prostate cancer datasets. This enhancement enables the system to leverage weakly supervised signals for improved nuclei-level analysis.

Highlights

• New Meta-Architecture for Weakly Supervised Learning (WSL): Introduced a new NucleiWSLMetaArch class, implemented using PyTorch Lightning, to handle weakly supervised learning on nuclei graphs.
• New Training Configurations: Added two new configuration files for training the nuclei_wsl meta-architecture: one for prostate cancer MMCI data and another for Radboud PANDA data, both utilizing a Self Attention Transformer model.
• PyTorch Lightning Integration: The new meta-architecture leverages PyTorch Lightning for structured training, validation, and testing, including advanced optimizer and learning rate scheduling (warmup and cosine annealing).
• Dependency Update: Added lightning>=2.6.0 to the project's dependencies to support the new PyTorch Lightning implementation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces the NucleiWSLMetaArch PyTorch Lightning module and associated Hydra configurations for nucleus-level weakly supervised learning on prostate cancer datasets. The implementation includes training and evaluation logic, metric logging, and an optimizer setup with a linear warmup and cosine annealing scheduler. Feedback focuses on ensuring training/inference consistency by applying block masking during training and improving reproducibility by registering the best validation loss as a persistent buffer for checkpointing.

[gemini-code-assist]

To ensure reproducibility and correct behavior when resuming training from a checkpoint, best_val_loss should be registered as a buffer. This ensures its value is persisted in the model's state dict and not reset to infinity upon restart.

Suggested change

	self.best_val_loss = float("inf")
	self.register_buffer("best_val_loss", torch.tensor(float("inf")))

References

1. Reproducibility is paramount. Ensure state that affects training logic is persisted. ^(link)

[gemini-code-assist]

If best_val_loss is registered as a buffer (tensor), it should be updated in-place or with a tensor to maintain its status as a buffer.

Suggested change

	self.best_val_loss = val_loss
	self.best_val_loss.fill_(val_loss)