Replace average_on_batch by average_on_trip

[michaelteboul]

Description of the goal of the PR

Description:

Changes this PR introduces (fill it before implementation)

• : Change 1
• : Change 2

Checklist before requesting a review

• I have commented my code, particularly in hard-to-understand areas
• I have typed my code
• I have created / updated the docstrings
• I have updated the README, if relevant
• I have updated the requirements files if a new package is used
• I have tested my code
• The CI pipeline passes
• I have performed a self-review of my code

[gemini-code-assist]

Summary of Changes

Hello @michaelteboul, 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 refines the metric calculation within the choice_learn library by improving the clarity and correctness of how metrics are averaged. It renames a key averaging parameter to average_on_trip and updates the underlying TensorFlow implementation to ensure that Negative Log Likelihood, Mean Reciprocal Rank, and Hit Rate are accurately aggregated on a per-trip basis, which is essential for consistent and meaningful evaluation in relevant contexts.

Highlights

• Parameter Renaming: The parameter average_on_batch has been consistently renamed to average_on_trip across the NegativeLogLikelihood, MRR (Mean Reciprocal Rank), and HitRate metric classes. This change clarifies the intent of the averaging mechanism, indicating that metrics are to be averaged per 'trip' rather than per generic 'batch'.
• Averaging Logic Update: The internal logic within the update_state methods of NegativeLogLikelihood, MRR, and HitRate has been updated. When average_on_trip is enabled, the metrics now correctly use tf.unique and tf.math.unsorted_segment_mean to compute averages based on unique 'trips' identified by the batch argument, ensuring accurate aggregation per trip.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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 metric classes by renaming average_on_batch to average_on_trip and implementing a more efficient and correct averaging logic using tf.math.unsorted_segment_mean. This is a significant improvement over the previous implementation, which appeared to be buggy and inefficient.

My review includes a few minor suggestions to improve code clarity and consistency. Most importantly, I've noticed that there are no unit tests for the new average_on_trip functionality. I strongly recommend adding tests to verify the correctness of the new implementation for all three metric classes (NegativeLogLikeliHood, MRR, and HitRate).

[gemini-code-assist]

The docstring on the following lines (44-45) for this parameter still refers to 'batch'. Please update it to 'trip' for consistency with the parameter name change.

[gemini-code-assist]

This commented-out code should be removed to improve code clarity.

[gemini-code-assist]

For consistency with the if branch and for better readability, consider using assign_add here. It would also be good to apply the same change to the update of self.n_evals on the next line for consistency.

Suggested change

	self.mrr.assign(self.mrr + tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32)))
	self.mrr.assign_add(tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32)))

[github-actions]

Coverage Report for Python 3.9

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	62	1	98%	283
choice_learn/basket_models
__init__.py	4	0	100%
alea_carta.py	148	22	85%	86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
base_basket_model.py	235	67	71%	111–112, 123, 141, 185, 255, 284–299, 377, 485, 585–587, 676, 762, 772, 822–830, 891–921, 931–953, 964–967, 991–1014
basic_attention_model.py	89	4	96%	424, 427, 433, 440
self_attention_model.py	133	9	93%	71, 73, 75, 450–454, 651
shopper.py	184	9	95%	130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
__init__.py	2	0	100%
basket_dataset.py	190	50	74%	74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 774–849
preprocessing.py	94	78	17%	43–45, 128–364
choice_learn/basket_models/datasets
__init__.py	3	0	100%
bakery.py	38	3	92%	47, 51, 61
synthetic_dataset.py	81	6	93%	62, 194–199, 247
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	649	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	241	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	400	5	99%	42–43, 153–154, 714
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	325	27	92%	145, 187, 289, 297, 303, 312, 352, 356–357, 362, 391, 395–396, 413, 426, 434, 475–476, 485–486, 587, 589, 605, 609, 611, 734–735
baseline_models.py	49	0	100%
conditional_logit.py	269	26	90%	49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
halo_mnl.py	124	18	85%	186, 341, 360, 364–380
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	236	236	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
metrics.py	90	51	43%	74, 94–99, 128–132, 149–172, 182, 196–205, 217–240, 250
TOTAL	5639	927	84%

Tests	Skipped	Failures	Errors	Time
222	0 💤	5 ❌	0 🔥	6m 24s ⏱️

[github-actions]

Coverage Report for Python 3.10

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	62	1	98%	283
choice_learn/basket_models
__init__.py	4	0	100%
alea_carta.py	148	22	85%	86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
base_basket_model.py	235	67	71%	111–112, 123, 141, 185, 255, 284–299, 377, 485, 585–587, 676, 762, 772, 822–830, 891–921, 931–953, 964–967, 991–1014
basic_attention_model.py	89	4	96%	424, 427, 433, 440
self_attention_model.py	133	9	93%	71, 73, 75, 450–454, 651
shopper.py	184	9	95%	130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
__init__.py	2	0	100%
basket_dataset.py	190	50	74%	74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 774–849
preprocessing.py	94	78	17%	43–45, 128–364
choice_learn/basket_models/datasets
__init__.py	3	0	100%
bakery.py	38	3	92%	47, 51, 61
synthetic_dataset.py	81	6	93%	62, 194–199, 247
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	649	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	241	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	400	5	99%	42–43, 153–154, 714
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	325	27	92%	145, 187, 289, 297, 303, 312, 352, 356–357, 362, 391, 395–396, 413, 426, 434, 475–476, 485–486, 587, 589, 605, 609, 611, 734–735
baseline_models.py	49	0	100%
conditional_logit.py	269	26	90%	49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
halo_mnl.py	124	18	85%	153, 172, 176–192, 374
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
metrics.py	90	51	43%	74, 94–99, 128–132, 149–172, 182, 196–205, 217–240, 250
TOTAL	5641	929	84%

Tests	Skipped	Failures	Errors	Time
222	0 💤	5 ❌	0 🔥	6m 51s ⏱️

[github-actions]

Coverage Report for Python 3.11

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	62	1	98%	283
choice_learn/basket_models
__init__.py	4	0	100%
alea_carta.py	148	22	85%	86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
base_basket_model.py	235	67	71%	111–112, 123, 141, 185, 255, 284–299, 377, 485, 585–587, 676, 762, 772, 822–830, 891–921, 931–953, 964–967, 991–1014
basic_attention_model.py	89	4	96%	424, 427, 433, 440
self_attention_model.py	133	9	93%	71, 73, 75, 450–454, 651
shopper.py	184	9	95%	130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
__init__.py	2	0	100%
basket_dataset.py	190	50	74%	74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 774–849
preprocessing.py	94	78	17%	43–45, 128–364
choice_learn/basket_models/datasets
__init__.py	3	0	100%
bakery.py	38	3	92%	47, 51, 61
synthetic_dataset.py	81	6	93%	62, 194–199, 247
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	649	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	241	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	400	5	99%	42–43, 153–154, 714
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	325	27	92%	145, 187, 289, 297, 303, 312, 352, 356–357, 362, 391, 395–396, 413, 426, 434, 475–476, 485–486, 587, 589, 605, 609, 611, 734–735
baseline_models.py	49	0	100%
conditional_logit.py	269	26	90%	49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
halo_mnl.py	124	18	85%	186, 341, 360, 364–380
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
metrics.py	90	51	43%	74, 94–99, 128–132, 149–172, 182, 196–205, 217–240, 250
TOTAL	5641	929	84%

Tests	Skipped	Failures	Errors	Time
222	0 💤	5 ❌	0 🔥	7m 5s ⏱️

[github-actions]

Coverage Report for Python 3.12

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	62	1	98%	283
choice_learn/basket_models
__init__.py	4	0	100%
alea_carta.py	148	22	85%	86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
base_basket_model.py	235	67	71%	111–112, 123, 141, 185, 255, 284–299, 377, 485, 585–587, 676, 762, 772, 822–830, 891–921, 931–953, 964–967, 991–1014
basic_attention_model.py	89	4	96%	424, 427, 433, 440
self_attention_model.py	133	9	93%	71, 73, 75, 450–454, 651
shopper.py	184	9	95%	130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
__init__.py	2	0	100%
basket_dataset.py	190	50	74%	74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 774–849
preprocessing.py	94	78	17%	43–45, 128–364
choice_learn/basket_models/datasets
__init__.py	3	0	100%
bakery.py	38	3	92%	47, 53, 61
synthetic_dataset.py	81	6	93%	62, 194–199, 247
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	649	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	241	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	400	5	99%	42–43, 153–154, 714
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	325	27	92%	145, 187, 289, 297, 303, 312, 352, 356–357, 362, 391, 395–396, 413, 426, 434, 475–476, 485–486, 587, 589, 605, 609, 611, 734–735
baseline_models.py	49	0	100%
conditional_logit.py	269	26	90%	49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
halo_mnl.py	124	18	85%	153, 172, 176–192, 374
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
metrics.py	90	51	43%	74, 94–99, 128–132, 149–172, 182, 196–205, 217–240, 250
TOTAL	5641	929	84%

Tests	Skipped	Failures	Errors	Time
222	0 💤	5 ❌	0 🔥	7m 28s ⏱️