feat: modeling

[MarcelKrkos]

LightningModule for tile-level supervised methodology PC-109

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for embedding-based classification models with configurable decode heads.
  • Introduced Virchow2 model configuration for multi-class classification tasks.

• Refactor

  • Refactored model training architecture for improved modularity and flexibility.

• Chores

  • Added seaborn to development dependencies.

[coderabbitai]

Warning

Rate limit exceeded

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

You’ve run out of usage credits. Purchase more in the billing tab.

⌛ 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: ca3b4445-b91e-491e-abf3-ebebee697015

📥 Commits

Reviewing files that changed from the base of the PR and between 6048953 and b219b87.

📒 Files selected for processing (2)

• ml/base.py
• ml/modeling/decode_head/base.py

📝 Walkthrough

Walkthrough

This PR refactors the model architecture from a single-class MetaArch orchestrator to a flexible base class pattern. A new GleasonModel Lightning base class provides training, validation, testing, and metric collection. EmbeddingGleasonModel concretely implements this base by delegating forward passes to a pluggable decode_head module. The Classifier abstraction enables reusable decode head implementations. Hydra configuration now controls model instantiation entirely.

Changes

Model architecture refactoring: MetaArch → GleasonModel

Layer / File(s)	Summary
Classifier abstraction for decode heads ml/modeling/decode_head/base.py, ml/modeling/decode_head/embedding_classifier.py, ml/modeling/decode_head/__init__.py	Abstract Classifier base class with dropout and linear projection, concrete EmbeddingClassifier validates 2D input and applies projection, module exports both via __all__.
GleasonModel Lightning base class ml/base.py	New abstract base for multiclass models built on LightningModule. Defines CrossEntropyLoss criterion, multiclass MetricCollection (AUROC, accuracy, precision, recall, F1, specificity, NPV), and implements training_step, validation_step, test_step, predict_step, and configure_optimizers with AdamW.
EmbeddingGleasonModel concrete implementation ml/embedding_model.py	Subclasses GleasonModel, stores and delegates forward computation to a pluggable decode_head nn.Module.
Configuration and entry point updates configs/model/decode_head/embedding_classifier.yaml, configs/model/virchow2.yaml, ml/__main__.py	New Hydra configs define embedding classifier (in/out features, dropout) and virchow2 model (target, 3 classes, decode head wiring). Remove MetaArch import and override from main so model class is determined by config.
Development dependencies update pyproject.toml	Add seaborn>=0.13.2 to dev dependency group.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested reviewers

• ejdam87

🐰 A Lightning framework springs forth so bright,
With Gleason models standing tall in the light,
No MetaArch chains to bind the design,
Classifiers pluggable, metrics align,
Configuration orchestrates—a refactor divine!

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title 'feat: modeling' is overly vague and does not clearly convey what the main change is, using a non-descriptive term 'modeling' without explaining the specific feature or architectural change being introduced.	Provide a more specific and descriptive title that explains the main architectural change, such as 'feat: add tile-level embedding model with classification head' or 'feat: refactor model architecture with embedding classifier.'

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/model

---

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 refined modeling framework for tile-level supervised learning. By replacing the previous monolithic architecture with a more flexible, base-class-driven design, the changes improve code maintainability and allow for easier experimentation with different decoding heads. The update also includes necessary configuration files and dependency updates to support these new architectural components.

Highlights

• Refactored Model Architecture: Replaced the generic MetaArch with a structured GleasonModel base class and specific implementations, improving modularity for tile-level supervision.
• Added Decoding Head Components: Introduced a new modular decode_head architecture, including a base Classifier class and an EmbeddingClassifier implementation.
• Configuration Updates: Added new Hydra configurations for the Virchow2 model and the embedding classifier to support the new modeling structure.
• Dependency Management: Added seaborn to the development dependencies in pyproject.toml and updated the lock file.

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 the 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 counterproductive. 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 refactors the model architecture by replacing the MetaArch class with a modular GleasonModel base class and introducing a new decode_head module structure. It also updates Hydra configurations and adds seaborn as a dependency. Feedback suggests improving MLflow experiment tracking by logging training loss at the epoch level and disabling step-level logging for noisy metrics like AUROC and F1-score to ensure cleaner visualizations.

[MarcelKrkos]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the model architecture by replacing the MetaArch class with a modular system based on a new GleasonModel base class and a specific EmbeddingGleasonModel implementation. It introduces a Classifier abstraction for decode heads, including an EmbeddingClassifier, and updates the Hydra configurations to support these changes. Additionally, seaborn is added as a development dependency. Feedback focuses on improving experiment reproducibility by calling self.save_hyperparameters() in the model constructors, ensuring that configuration parameters are preserved in checkpoints while avoiding serialization issues with complex sub-modules.