(
self,
data_root,
dataset,
batch_size,
parallel_loading=True,
num_workers=256,
logger=None,
cache_dir="~/cache/images_lmdb_sharded",
lmdb_num_shards=32,
lmdb_map_size_per_shard=1024 * 1024 * 1024 * 1024, # 1TB per shard
lmdb_readonly=False,
lmdb_lock=False,
resize_image=True,
resize_image_size=1024,
use_image=True,
use_lmdb_cache=True, # Whether to cache images in LMDB; False = read from disk directly
adjust_batch_size=True, # Whether to adjust batch_size based on turns
checkpoint_manager=None, # Checkpoint manager for item-level resume
**kwargs, # Accept but ignore deprecated parameters like pre_load_img
)
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | data_root, |
| 40 | dataset, |
| 41 | batch_size, |
| 42 | parallel_loading=True, |
| 43 | num_workers=256, |
| 44 | logger=None, |
| 45 | cache_dir="~/cache/images_lmdb_sharded", |
| 46 | lmdb_num_shards=32, |
| 47 | lmdb_map_size_per_shard=1024 * 1024 * 1024 * 1024, # 1TB per shard |
| 48 | lmdb_readonly=False, |
| 49 | lmdb_lock=False, |
| 50 | resize_image=True, |
| 51 | resize_image_size=1024, |
| 52 | use_image=True, |
| 53 | use_lmdb_cache=True, # Whether to cache images in LMDB; False = read from disk directly |
| 54 | adjust_batch_size=True, # Whether to adjust batch_size based on turns |
| 55 | checkpoint_manager=None, # Checkpoint manager for item-level resume |
| 56 | **kwargs, # Accept but ignore deprecated parameters like pre_load_img |
| 57 | ): |
| 58 | self.data_root = data_root |
| 59 | self.batch_size = batch_size |
| 60 | self.original_config = None |
| 61 | self.parallel_loading = parallel_loading |
| 62 | self.num_workers = num_workers |
| 63 | self.logger = logger |
| 64 | self.resize_image = resize_image |
| 65 | self.resize_image_size = resize_image_size |
| 66 | self.use_image = use_image |
| 67 | self.use_lmdb_cache = use_lmdb_cache |
| 68 | self.adjust_batch_size = adjust_batch_size |
| 69 | self.checkpoint_manager = checkpoint_manager |
| 70 | self.name = "" |
| 71 | |
| 72 | # Store LMDB config for lazy initialization |
| 73 | self._cache_dir = os.path.expanduser(cache_dir) if cache_dir else None |
| 74 | self._lmdb_num_shards = lmdb_num_shards |
| 75 | self._lmdb_map_size_per_shard = lmdb_map_size_per_shard |
| 76 | self._lmdb_readonly = lmdb_readonly |
| 77 | self._lmdb_lock = lmdb_lock |
| 78 | self.sharded_lmdb = None |
| 79 | |
| 80 | # Resolve dataset path and load data (supports JSON/JSONL formats) |
| 81 | data_file = self._resolve_data_path(dataset, data_root) |
| 82 | self.data_file = data_file |
| 83 | self.all_datas = FormatRegistry.load( |
| 84 | data_file, add_source_file=True, remove_rejected=True |
| 85 | ) |
| 86 | self.total_items = len(self.all_datas) |
| 87 | |
| 88 | # Apply checkpoint: skip already processed items |
| 89 | self.start_idx = 0 |
| 90 | if self.checkpoint_manager: |
| 91 | self.start_idx = self.checkpoint_manager.get_start_index( |
| 92 | data_file, self.total_items |
| 93 | ) |
| 94 | if self.start_idx == -1: |
nothing calls this directly
no test coverage detected