Add adjusted_filter_batch (#21260)

This commit is contained in:
Aurick Qiao
2026-03-26 10:59:05 +08:00
committed by GitHub
parent 53c1d8e963
commit a34e9ed64a
@@ -191,6 +191,12 @@ class SamplingBatchInfo:
def adjusted_merge_batch(self, other: "SamplingBatchInfo"): def adjusted_merge_batch(self, other: "SamplingBatchInfo"):
pass pass
# placeholder for override
def adjusted_filter_batch(
self, keep_indices: List[int], keep_indices_device: torch.Tensor
):
pass
def __len__(self): def __len__(self):
return len(self.temperatures) return len(self.temperatures)
@@ -267,6 +273,8 @@ class SamplingBatchInfo:
if self.logit_bias is not None: if self.logit_bias is not None:
self.logit_bias = self.logit_bias[keep_indices_device] self.logit_bias = self.logit_bias[keep_indices_device]
self.adjusted_filter_batch(keep_indices, keep_indices_device)
def _filter_batch_custom_logit_processor( def _filter_batch_custom_logit_processor(
self, keep_indices: List[int], keep_indices_device: torch.Tensor self, keep_indices: List[int], keep_indices_device: torch.Tensor
): ):