renamed lib

Author: finbarrtimbers

open_instruct/data_loader.py

@@ -34,7 +34,7 @@
 from transformers import PreTrainedTokenizer
 
 from open_instruct import data_types, padding_free_collator, utils
-from open_instruct import replay_buffer as replay_buffer_mod
+from open_instruct import replay_buffer as replay_buffer_lib
 from open_instruct.data_types import EnvConfig, EnvConfigEntry
 from open_instruct.dataset_transformation import (
     ENV_CONFIG_KEY,
@@ -487,14 +487,14 @@ class StreamingDataLoaderConfig:
     save_traces: bool = False
     rollouts_save_path: str = "/weka/oe-adapt-default/allennlp/deletable_rollouts/"
 
-    replay_buffer: replay_buffer_mod.ReplayBufferConfig = field(default_factory=replay_buffer_mod.ReplayBufferConfig)
+    replay_buffer: replay_buffer_lib.ReplayBufferConfig = field(default_factory=replay_buffer_lib.ReplayBufferConfig)
 
     # Computed at post_init
     max_possible_score: float = 1.0
 
     def __post_init__(self):
         if isinstance(self.replay_buffer, dict):
-            self.replay_buffer = replay_buffer_mod.ReplayBufferConfig(**self.replay_buffer)
+            self.replay_buffer = replay_buffer_lib.ReplayBufferConfig(**self.replay_buffer)
         assert self.pack_length >= self.max_prompt_token_length + self.response_length, (
             "The `pack_length` needs to be greater than the sum of `max_prompt_token_length` and `response_length`!"
         )
@@ -745,7 +745,7 @@ def process_single_result(
     replenish_prompts: bool,
     param_prompt_Q: ray_queue.Queue | None,
     base_env_config: EnvConfig,
-) -> replay_buffer_mod.ProcessedResult | None:
+) -> replay_buffer_lib.ProcessedResult | None:
     assert result.index is not None
     assert result.logprobs is not None
     assert result.reward_scores is not None
@@ -793,7 +793,7 @@ def process_single_result(
     if filter_zero_std_samples and np.std(result.reward_scores) == 0:
         return None
 
-    return replay_buffer_mod.ProcessedResult(
+    return replay_buffer_lib.ProcessedResult(
         result=result,
         queries=repeat_each([query], generation_config.n),
         ground_truths=repeat_each([ground_truth], generation_config.n),
@@ -808,7 +808,7 @@ def process_single_result(
 
 
 def combine_processed_results(
-    processed_results: list[replay_buffer_mod.ProcessedResult],
+    processed_results: list[replay_buffer_lib.ProcessedResult],
     generation_config: vllm.SamplingParams,
     actor_manager=None,
 ) -> tuple[data_types.GenerationResult, Batch, dict, BatchStatistics] | tuple[None, None, None, None]:
@@ -975,7 +975,7 @@ def accumulate_inference_batches(
             "replenish_prompts requires param_prompt_Q and iter_dataloader and dataset"
         )
 
-    processed_results: list[replay_buffer_mod.ProcessedResult] = []
+    processed_results: list[replay_buffer_lib.ProcessedResult] = []
     total_filtered_prompts = 0
     filtered_prompt_zero = 0
     filtered_prompt_solved = 0
@@ -1239,12 +1239,12 @@ def __init__(
         self.metadata_saved = False
 
         rb = config.replay_buffer
-        self._table = replay_buffer_mod.Table(
+        self._table = replay_buffer_lib.Table(
             max_size=rb.capacity or global_batch_size,
-            sampler=replay_buffer_mod.make_selector(rb.sampler),
-            remover=replay_buffer_mod.make_selector(rb.remover),
+            sampler=replay_buffer_lib.make_selector(rb.sampler),
+            remover=replay_buffer_lib.make_selector(rb.remover),
             max_times_sampled=rb.max_times_sampled,
-            rate_limiter=replay_buffer_mod.MinSize(rb.min_size or global_batch_size),
+            rate_limiter=replay_buffer_lib.MinSize(rb.min_size or global_batch_size),
         )
 
         if initial_state is not None: