()
| 129 | |
| 130 | |
| 131 | def test_sampler(): |
| 132 | batch_size = 32 |
| 133 | vocab_size = 1024 |
| 134 | min_seq_len = 1 |
| 135 | max_seq_len = 1024 |
| 136 | |
| 137 | sampler = Sampler(get_fd_config(batch_size)) |
| 138 | logits = _create_fake_logits(batch_size, vocab_size) |
| 139 | sampling_metadata = _create_default_sampling_metadata(batch_size, min_seq_len, max_seq_len) |
| 140 | next_tokens = sampler(logits, sampling_metadata) |
| 141 | print(next_tokens) |
| 142 | |
| 143 | |
| 144 | def get_baseline_logprobs(logits, sampling_metadata, logprobs_mode, token_ids): |
no test coverage detected