quaterion.dataset.train_collator module

class TrainCollator(pre_collate_fn: Callable, encoder_collates: Dict[str, CollateFnType], meta_extractors: Dict[str, Callable[[List[Any]], List[dict]]])

Bases: object

Functional object, that aggregates all required information for performing collate on train batches.

Note

Should be serializable for sending among worker processes.

Parameters:

• pre_collate_fn – function to split origin batch into ids, features and labels. Ids are means to keep track of repeatable usage of the same elements. Features are commonly encoders input. Labels usually allow distinguishing positive and negative samples.

• encoder_collates – mapping of encoder name to its collate function

pre_encoder_collate(features: List[Any], ids: List[int] | None = None, encoder_name: str | None = None)

Default implementation of per-encoder batch preparation, might be overridden

process_meta(meta: Dict[str, List]) → Any