trinity.buffer.selector package

Submodules

• trinity.buffer.selector.difficulty_estimator module
  • BaseBetaPREstimator
    • BaseBetaPREstimator.__init__()
    • BaseBetaPREstimator.n
    • BaseBetaPREstimator.m
    • BaseBetaPREstimator.lamb
    • BaseBetaPREstimator.rho
    • BaseBetaPREstimator.alphas
    • BaseBetaPREstimator.betas
    • BaseBetaPREstimator.set()
    • BaseBetaPREstimator.update()
    • BaseBetaPREstimator.predict_pr()
    • BaseBetaPREstimator.equivalent_count()
  • InterpolationBetaPREstimator
    • InterpolationBetaPREstimator.__init__()
    • InterpolationBetaPREstimator.update()
• trinity.buffer.selector.selector module
  • BaseSelector
    • BaseSelector.__init__()
    • BaseSelector.get_indices()
    • BaseSelector.update()
    • BaseSelector.state_dict()
    • BaseSelector.load_state_dict()
  • SequentialSelector
    • SequentialSelector.__init__()
    • SequentialSelector.get_indices()
    • SequentialSelector.update()
    • SequentialSelector.state_dict()
    • SequentialSelector.load_state_dict()
  • ShuffleSelector
    • ShuffleSelector.__init__()
    • ShuffleSelector.get_indices()
    • ShuffleSelector.update()
    • ShuffleSelector.state_dict()
    • ShuffleSelector.load_state_dict()
  • RandomSelector
    • RandomSelector.__init__()
    • RandomSelector.get_indices()
    • RandomSelector.update()
    • RandomSelector.state_dict()
    • RandomSelector.load_state_dict()
  • OfflineEasy2HardSelector
    • OfflineEasy2HardSelector.__init__()
    • OfflineEasy2HardSelector.update()
    • OfflineEasy2HardSelector.get_indices()
    • OfflineEasy2HardSelector.state_dict()
    • OfflineEasy2HardSelector.load_state_dict()
  • DifficultyBasedSelector
    • DifficultyBasedSelector.__init__()
    • DifficultyBasedSelector.build_diff_estimator()
    • DifficultyBasedSelector.update()
    • DifficultyBasedSelector.get_scores()
    • DifficultyBasedSelector.get_indices()
    • DifficultyBasedSelector.state_dict()
    • DifficultyBasedSelector.load_state_dict()

Module contents

class trinity.buffer.selector.BaseSelector(data_source: _HFBatchReader, config: TaskSelectorConfig)

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)

get_indices(batch_size: int, return_extra_info: bool = False) → List[int]

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

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

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

Restore selector state from a saved dictionary.

Parameters:

state_dict (Dict) – Output from state_dict()