| |
| |
|
|
| """Host-only state contracts for TT sampled-token log-probability mode.""" |
|
|
| from models.common.sampling.tt_log_probs import LogProbsCalculator |
|
|
|
|
| def _calculator(batch_size=4): |
| calculator = object.__new__(LogProbsCalculator) |
| calculator.batch_size = batch_size |
| calculator.logprobs_enabled = [False] * batch_size |
| calculator.num_logprobs = [0] * batch_size |
| calculator.enable_log_probs = False |
| calculator.topk_logprobs_needed = False |
| return calculator |
|
|
|
|
| def test_all_false_tuple_keeps_log_probs_disabled(): |
| calculator = _calculator() |
|
|
| calculator.set_log_probs_mode((False, False, False, False), num_logprobs=(0, 0, 0, 0)) |
|
|
| assert calculator.logprobs_enabled == [False, False, False, False] |
| assert calculator.num_logprobs == [0, 0, 0, 0] |
| assert calculator.enable_log_probs is False |
| assert calculator.topk_logprobs_needed is False |
|
|
|
|
| def test_tuple_log_probs_and_counts_map_per_user_exactly(): |
| calculator = _calculator() |
|
|
| calculator.set_log_probs_mode((True, False, True, False), num_logprobs=(0, 4, 7, 2)) |
|
|
| assert calculator.logprobs_enabled == [True, False, True, False] |
| assert calculator.num_logprobs == [0, 4, 7, 2] |
| assert calculator.enable_log_probs is True |
| assert calculator.topk_logprobs_needed is True |
|
|
|
|
| def test_tuple_partial_update_respects_empty_slots(): |
| calculator = _calculator(batch_size=6) |
|
|
| calculator.set_log_probs_mode((True, False), num_logprobs=(10, 15), empty_slots=[2, 5]) |
|
|
| assert calculator.logprobs_enabled == [False, False, True, False, False, False] |
| assert calculator.num_logprobs == [0, 0, 10, 0, 0, 15] |
| assert calculator.enable_log_probs is True |
|
|
|
|
| def test_existing_scalar_and_list_modes_remain_unchanged(): |
| calculator = _calculator() |
|
|
| calculator.set_log_probs_mode(True, num_logprobs=5) |
| assert calculator.logprobs_enabled == [True, True, True, True] |
| assert calculator.num_logprobs == [5, 5, 5, 5] |
|
|
| calculator.set_log_probs_mode([True, False, False, True], num_logprobs=[1, 2, 3, 4]) |
| assert calculator.logprobs_enabled == [True, False, False, True] |
| assert calculator.num_logprobs == [1, 2, 3, 4] |
| assert calculator.enable_log_probs is True |
| assert calculator.topk_logprobs_needed is True |
|
|
