Add reorder cache for beam search

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
+- Added _reorder_cache function to enable beam search when generating with HF
 
 ### Changed
 

hf_olmo/modeling_olmo.py

@@ -133,8 +133,15 @@ def prepare_inputs_for_generation(
     # def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]:
     #     pass
     #
-    # def _reorder_cache(self, past_key_values, beam_idx):
-    #     pass
+
+    def _reorder_cache(self, past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
+    ) -> Tuple[Tuple[torch.Tensor]]:
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
 
     def get_input_embeddings(self) -> torch.nn.Module:
         return self.model.transformer.wte

scripts/add_code_eval.py

@@ -1,6 +1,7 @@
 """
 Script to create perplexity eval datasets for code.
 """
+
 import os
 
 import pandas as pd

scripts/init_config.py

@@ -1,6 +1,7 @@
 """
 Run this to initialize a new training config to a file.
 """
+
 import logging
 import sys
 from pathlib import Path

scripts/show_model_size.py

@@ -7,6 +7,7 @@
 python scripts/show_model_size.py train_config.yaml
 ```
 """
+
 import logging
 import sys