Spaces:
Running
on
A100
Running
on
A100
File size: 3,659 Bytes
174ae06 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# Copyright (c) 2025 NVIDIA CORPORATION.
# Licensed under the MIT license.
# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
# LICENSE is in incl_licenses directory.
import random
from typing import Any, Dict, List
import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
from llava.mm_utils import dynamic_process_images_and_prompt, dynamic_s2_process_images_and_prompt, process_images
from llava.train.args import DataArguments
from llava.utils.logging import logger
from llava.utils.media import extract_media
from llava.utils.tokenizer import preprocess_conversation
__all__ = ["BaseDataset"]
def _process_speech(speech: List[Any], data_args: DataArguments) -> torch.Tensor:
return torch.tensor(speech)
def _process_sound(sound: List[Any], data_args: DataArguments) -> torch.Tensor:
return torch.tensor(sound)
def _process_sound_masks(sound_masks: List[Any], data_args: DataArguments) -> torch.Tensor:
return torch.tensor(sound_masks)
class BaseDataset(Dataset):
def __init__(
self,
tokenizer: PreTrainedTokenizer,
data_args: DataArguments,
no_system_prompt: bool = False,
**kwargs: Any,
) -> None:
super().__init__()
self.tokenizer = tokenizer
self.data_args = data_args
self.no_system_prompt = no_system_prompt
self.instances = []
self.enable_dynamic_res = False
self.enable_dynamic_res_s2 = False
# global_batch_size: int,
self.global_batch_size = kwargs.get("global_batch_size", 1)
# by default, dataset cls will resample on failure
self.resample_on_failure = kwargs.get("resample_on_failure", True)
# by default, dataset cls will resample on failure
self.resample_on_failure = kwargs.get("resample_on_failure", True)
def process(self, instance: Dict[str, Any]) -> List[Dict[str, Any]]:
raise NotImplementedError
def __getitem__(self, index: int) -> Dict[str, Any]:
instance = self.instances[index]
try:
# Process instance to conversation
conversation = self.process(instance)
# Extract media from conversation
media, media_meta = extract_media(conversation, self.data_args)
if "speech" in media:
processed_speech = _process_speech(media["speech"], self.data_args)
if "sound" in media:
processed_sound = _process_sound(media["sound"], self.data_args)
processed_sound_feature_masks = _process_sound_masks(media_meta["sound_feature_masks"], self.data_args)
processed_sound_embed_masks = _process_sound_masks(media_meta["sound_embed_masks"], self.data_args)
# Prepare "input_ids" and "labels" for training
data = preprocess_conversation(conversation, self.tokenizer, no_system_prompt=self.no_system_prompt)
if "speech" in media:
data["speech"] = processed_speech
if "sound" in media:
data["sound"] = processed_sound
data["sound_feature_masks"] = processed_sound_feature_masks
data["sound_embed_masks"] = processed_sound_embed_masks
except Exception as e:
if not self.resample_on_failure:
raise e
else:
logger.exception(f"Error processing instance '{instance}': '{e}'. Resampling.")
return self.__getitem__(random.randint(0, len(self.instances) - 1))
return data
def __len__(self) -> int:
return len(self.instances)
|