test_common_inference_params.py 396 Bytes
Newer Older
silencealiang's avatar
add  
silencealiang committed
1
from megatron.core.inference.sampling_params import SamplingParams
xingjinliang's avatar
xingjinliang committed
2
3


silencealiang's avatar
add  
silencealiang committed
4
class TestSamplingParams:
xingjinliang's avatar
xingjinliang committed
5
6

    def test_inference_params(self):
silencealiang's avatar
add  
silencealiang committed
7
        inference_parameters = SamplingParams()
xingjinliang's avatar
xingjinliang committed
8
9
10
11
        inference_parameters.add_attributes({"min_tokens": 45})
        assert (
            inference_parameters.min_tokens == 45
        ), f"min tokens not set correctly. it is {inference_parameters.min_tokens}"