LLM: update benchmark_util.py for beam search (#9126)

* update reorder_cache

* fix
This commit is contained in:
Ruonan Wang 2023-10-11 09:41:53 +08:00 committed by GitHub
parent e8c5645067
commit 1363e666fc

View file

@ -821,10 +821,7 @@ class BenchmarkWrapper:
return model_kwargs
def _reorder_cache(self, past_key_values, beam_idx):
raise NotImplementedError(
f"Make sure that a `_reorder_cache` function is correctly implemented in {self.__class__.__module__} to"
f" enable beam search for {self.__class__}"
)
return self.model._reorder_cache(past_key_values, beam_idx)
def _get_logits_warper(
self,