Class ExamplePreprocessor

olmo/preprocessing/multimodal_preprocessor.py:189–311

Preprocesses examples dictionaries as returned by our data loaders Includes loading the multi-modal data and formatting the text for the LLM

@dataclasses.dataclass
class ExamplePreprocessor:
    """Preprocesses examples dictionaries as returned by our data loaders

    Includes loading the multi-modal data and formatting the text for the LLM
    """
    formatter: DataFormatter
    preprocessor: MultimodalPreprocessor
    for_inference: bool = False
    is_training: bool = False
    include_image: bool = False

    def get_output_shapes(self) -> Dict[str, TensorSpec]:
        return self.preprocessor.get_output_shapes()

    @property
    def tokenizer(self):
        return self.preprocessor.text_preprocessor.tokenizer

    def __call__(self, example, rng=np.random):
        example = dict(example)
        image: Optional[np.ndarray] = None
        image_group: Optional[List[np.ndarray]] = None
        video: Optional[VideoFrames] = None
        if "image" in example:
            is_image_group = isinstance(example["image"], (list, tuple))
            try:
                if is_image_group:
                    image_group = [load_image(x) for x in example["image"]]
                else:
                    image = load_image(example["image"])
            except Exception as e:
                e.add_note(f"Could not load image: {example['image']}")
                raise e
            if not is_image_group:
                # So the formatter can know the height/weight of the video
                example["image"] = image
            image_to_video_metadata = None
        if "images" in example:
            example["images"] = [load_image(x) for x in example["images"]]

        if "video" in example:
            if isinstance(example["video"], VideoFrames):
                video_path = None
                video = example["video"]
            else:
                video_path = example["video"]
                try:
                    decode_method = None
                    if "metadata" in example and "decode_method" in example["metadata"]:
                        decode_method = example["metadata"]["decode_method"]
                    clip = None
                    if "metadata" in example and "clip_start_time" in example["metadata"]:
                        if example["metadata"]["clip_start_time"] is not None:
                            clip = (example["metadata"]["clip_start_time"], example["metadata"]["clip_end_time"])
                    subtitle = None
                    if 'subtitle' in example:
                        subtitle = example['subtitle']
                    sampler_overrides = {}