flashinfer.logits_processor.Temperature

class flashinfer.logits_processor.Temperature(**params: Any)

用于 logits 的温度缩放处理器。

通过除以温度值来缩放 logits。

TensorType.LOGITS -> TensorType.LOGITS

参数:

temperature (floattorch.Tensor, Runtime) – 用于缩放的温度值。必须为正数。可以是一个标量或每批次的张量。

示例

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Temperature, Sample
>>> torch.manual_seed(42)
>>> pipe = LogitsPipe([Temperature()])
>>> logits = torch.randn(2, 2, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614], [ -0.1721,  0.8491]], device='cuda:0')
>>> scaled_logits = pipe(logits, temperature=0.8)
>>> scaled_logits
tensor([[ 0.2425,  2.7017], [-0.2151,  1.0613]], device='cuda:0')
__init__(**params: Any)

温度处理器的构造函数。不需要编译时参数。

方法

__init__(**params)

温度处理器的构造函数。

legalize(input_type)

将处理器合法化为低级运算符列表。