trinity.buffer.selector package#

Submodules#

Module contents#

class trinity.buffer.selector.BaseSelector(data_source: _HFBatchReader, config: TaskSelectorConfig)[source]#

Bases: object

Abstract base class defining the interface for custom data selection strategies.

A selector determines which samples (by index) are selected from the dataset during training. It enables flexible sampling beyond simple sequential or random access, supporting active learning, curriculum learning, or difficulty-based sampling in the future.

Subclasses must implement:
  • get_indices: returns list of indices for next batch

  • update: updates internal state using feedback (e.g., loss values, mean rewards, etc.)

  • state_dict / load_state_dict: for saving/loading selector state (checkpointing)

__init__(data_source: _HFBatchReader, config: TaskSelectorConfig)[source]#
get_indices(batch_size: int, return_extra_info: bool = False) List[int][source]#

Select a batch of sample indices from the dataset.

Parameters:
  • batch_size (int) – Number of indices to return

  • return_extra_info (bool) – If True, may return additional metadata (future use)

Returns:

Selected indices into the dataset

Return type:

List[int]

update(indices: List[int], values: List[float]) None[source]#

Update internal state based on feedback (e.g., model loss, accuracy).

This allows adaptive selectors (like hard example mining) to learn over time.

Parameters:
  • indices (List[int]) – Previously selected indices

  • values (List[float]) – Feedback values corresponding to those indices

state_dict() Dict[source]#

Return serializable state of the selector for checkpointing.

Returns:

State information (e.g., current position, etc.)

Return type:

Dict

load_state_dict(state_dict: Dict) None[source]#

Restore selector state from a saved dictionary.

Parameters:

state_dict (Dict) – Output from state_dict()