Update 'balance probabilities test' and 'get_embeddings' test

[bejaeger]

Issue

This fixes 3 tests for upcoming V2.5:

• The get_embeddings didn't consider thinking tokens. Change this so that "train_embeddings" only returns the embeddings for the data points
• Balance probabilities test was flawed. We don't really expect the probabilities to be uniform for any model. This introduces a new test for the functionality of balance probabilities, as well as a softer criterion for the classifier, to test whether the proba distribution is more uniform with balance-probabilities included.
• test_classifier_in_pipeline: This was using the balanced probabilities check as the previous test which didn't make too much sense because we only want to test whether the classifier works in a pipeline in this test.

[gemini-code-assist]

Code Review

This pull request addresses issues in two tests. The get_embeddings test is updated to correctly handle 'thinking tokens' by adjusting the slice for training embeddings. The balance_probabilities test is rewritten with a more robust criterion, checking for increased uniformity in probability distributions rather than expecting a perfectly uniform one. This change also involves refactoring the balancing logic into a new utility function, balance_probas_by_class_counts, and adding a corresponding unit test. The changes are logical and improve the test suite. I've provided a critical fix for a new test that would otherwise fail due to incorrect parameters, and a suggestion to improve the robustness of the new utility function.

[Copilot]

Pull Request Overview

This PR refactors the probability balancing logic by extracting it into a standalone utility function and improves the corresponding tests. The changes focus on making the code more maintainable and testable.

Key Changes:

• Extracted balance_probas_by_class_counts utility function from classifier-specific code
• Added dedicated unit test for the probability balancing function
• Improved test coverage for balanced probabilities with more realistic imbalanced dataset scenarios
• Fixed train embeddings extraction in transformer to account for thinking tokens offset

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/tabpfn/utils.py	Added new balance_probas_by_class_counts function to centralize probability balancing logic
src/tabpfn/classifier.py	Refactored to use the extracted utility function instead of inline implementation
tests/test_utils.py	Added unit test for the new balance_probas_by_class_counts function
tests/test_classifier_interface.py	Updated test to use synthetic imbalanced dataset and improved assertion logic
src/tabpfn/architectures/base/transformer.py	Fixed train embeddings extraction to properly skip thinking token rows

Comments suppressed due to low confidence (1)

src/tabpfn/architectures/base/transformer.py:574

• This assignment to 'train_encoder_out' is unnecessary as it is redefined before this value is used.

                train_encoder_out = self.encoder_compression_layer(

---

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

[gemini-code-assist]

Code Review

This pull request effectively addresses two testing issues. The get_embeddings logic is corrected to account for thinking tokens, and the balance_probabilities test is significantly improved to be more robust and meaningful. The refactoring of the balancing logic into a separate utility function with its own unit test is a great improvement for code organization and maintainability.

I've provided a few suggestions to enhance code conciseness, improve numerical stability, and ensure test reproducibility. Overall, these are solid changes.

[LeoGrin]

LGTM! Agree that the previous balance probabilty test was too strong. (the new one might be a bit weak but don't have a great solution for that)