data_juicer.ops.selector.image_ohem_selector module#

Select hard image examples using a user supplied per-sample loss function.

class data_juicer.ops.selector.image_ohem_selector.ImageOHEMSelector(*args, **kwargs)[source]#

Bases: Selector

Keep the highest-loss image samples (Online Hard Example Mining).

score_fn is supplied by the user and must have the signature score_fn(model, samples, images, device, **score_kwargs). It must return one numeric loss per sample in the batch. model_factory is optional and, when supplied, is called once as model_factory(**model_kwargs).

__init__(score_fn=None, model_factory=None, score_file: str = '', score_function: str = 'score_fn', model_function: str = 'model_factory', top_ratio: float | None = None, topk: int | None = None, batch_size: int = 8, image_key: str = 'images', image_bytes_key: str = 'image_bytes', loss_field: str = 'image_ohem_loss', device: str = 'auto', model_kwargs: dict | None = None, score_kwargs: dict | None = None, *args, **kwargs)[source]#

Base class that conducts selection in dataset-level.

Parameters:

  • text_key – the key name of field that stores sample texts to be processed

  • image_key – the key name of field that stores sample image list to be processed

  • audio_key – the key name of field that stores sample audio list to be processed

  • video_key – the key name of field that stores sample video list to be processed

  • image_bytes_key – the key name of field that stores sample image bytes list to be processed

  • query_key – the key name of field that stores sample queries

  • response_key – the key name of field that stores responses

  • history_key – the key name of field that stores history of queries and responses

process(dataset)[source]#

Dataset –> dataset.

Parameters:

dataset – input dataset

Returns:

selected dataset.