714 lines
23 KiB
Python
714 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
from abc import ABC, abstractmethod
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
from os.path import join as opj
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from einops import rearrange
|
|
from PIL import Image
|
|
from transformers import AutoTokenizer
|
|
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class PreprocessBatch:
|
|
"""
|
|
Batch information for dataset processing stages.
|
|
|
|
This class holds all the information about a video-caption or image-caption pair
|
|
as it moves through the processing pipeline. Fields are populated by different stages.
|
|
"""
|
|
# Raw metadata
|
|
path: str
|
|
cap: str | list[str]
|
|
resolution: dict | None = None
|
|
fps: float | None = None
|
|
duration: float | None = None
|
|
|
|
# Processed metadata
|
|
num_frames: int | None = None
|
|
sample_frame_index: list[int] | None = None
|
|
sample_num_frames: int | None = None
|
|
action_path: str | None = None
|
|
|
|
# Processed data
|
|
pixel_values: torch.Tensor | None = None
|
|
text: str | None = None
|
|
input_ids: torch.Tensor | None = None
|
|
cond_mask: torch.Tensor | None = None
|
|
|
|
@property
|
|
def is_video(self) -> bool:
|
|
"""Check if this is a video item."""
|
|
return self.path.endswith(".mp4")
|
|
|
|
@property
|
|
def is_image(self) -> bool:
|
|
"""Check if this is an image item."""
|
|
return self.path.endswith(".jpg")
|
|
|
|
|
|
class DatasetStage(ABC):
|
|
"""
|
|
Abstract base class for dataset processing stages.
|
|
|
|
Similar to PipelineStage but designed for dataset preprocessing operations.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""
|
|
Process the dataset batch.
|
|
|
|
Args:
|
|
batch: Dataset batch to process
|
|
**kwargs: Additional processing parameters
|
|
|
|
Returns:
|
|
Processed batch
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class DatasetFilterStage(ABC):
|
|
"""
|
|
Abstract base class for dataset filtering stages.
|
|
|
|
These stages can filter out items during metadata processing.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
|
|
"""
|
|
Check if batch should be kept.
|
|
|
|
Args:
|
|
batch: Dataset batch to check
|
|
**kwargs: Additional parameters
|
|
|
|
Returns:
|
|
True if batch should be kept, False otherwise
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""
|
|
Process the dataset batch (for non-filtering operations).
|
|
|
|
Args:
|
|
batch: Dataset batch to process
|
|
**kwargs: Additional processing parameters
|
|
|
|
Returns:
|
|
Processed batch
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class DataValidationStage(DatasetFilterStage):
|
|
"""Stage for validating data items."""
|
|
|
|
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
|
|
"""
|
|
Validate data item.
|
|
|
|
Args:
|
|
batch: Dataset batch to validate
|
|
|
|
Returns:
|
|
True if valid, False if invalid
|
|
"""
|
|
# Check for caption
|
|
if batch.cap is None:
|
|
return False
|
|
|
|
if batch.is_video:
|
|
# Validate video-specific fields
|
|
if batch.duration is None or batch.fps is None:
|
|
return False
|
|
elif not batch.is_image:
|
|
return False
|
|
|
|
return True
|
|
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""Process does nothing for validation - filtering is handled by should_keep."""
|
|
return batch
|
|
|
|
|
|
class FrameSamplingStage(DatasetFilterStage):
|
|
"""Stage for temporal frame sampling and indexing."""
|
|
|
|
def __init__(self,
|
|
num_frames: int,
|
|
train_fps: int,
|
|
speed_factor: int = 1,
|
|
video_length_tolerance_range: float = 5.0,
|
|
drop_short_ratio: float = 0.0,
|
|
seed: int = 42):
|
|
self.num_frames = num_frames
|
|
self.train_fps = train_fps
|
|
self.speed_factor = speed_factor
|
|
self.video_length_tolerance_range = video_length_tolerance_range
|
|
self.drop_short_ratio = drop_short_ratio
|
|
# Create a seeded random generator for deterministic sampling
|
|
self.rng = random.Random(seed)
|
|
|
|
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
|
|
"""
|
|
Check if video should be kept based on length constraints.
|
|
|
|
Args:
|
|
batch: Dataset batch
|
|
|
|
Returns:
|
|
True if should be kept, False otherwise
|
|
"""
|
|
if batch.is_image:
|
|
return True
|
|
|
|
if batch.duration is None or batch.fps is None:
|
|
return False
|
|
|
|
num_frames = math.ceil(batch.fps * batch.duration)
|
|
|
|
# Check if video is too long
|
|
if (num_frames / batch.fps > self.video_length_tolerance_range *
|
|
(self.num_frames / self.train_fps * self.speed_factor)):
|
|
return False
|
|
|
|
# Resample frame indices to check length
|
|
frame_interval = batch.fps / self.train_fps
|
|
start_frame_idx = 0
|
|
frame_indices = np.arange(start_frame_idx, num_frames,
|
|
frame_interval).astype(int)
|
|
|
|
# Filter short videos
|
|
return not (len(frame_indices) < self.num_frames
|
|
and self.rng.random() < self.drop_short_ratio)
|
|
|
|
def process(self,
|
|
batch: PreprocessBatch,
|
|
temporal_sample_fn=None,
|
|
**kwargs) -> PreprocessBatch:
|
|
"""
|
|
Process frame sampling for video data items.
|
|
|
|
Args:
|
|
batch: Dataset batch
|
|
temporal_sample_fn: Function for temporal sampling
|
|
|
|
Returns:
|
|
Updated batch with frame sampling info
|
|
"""
|
|
if batch.is_image:
|
|
# For images, just add sample info
|
|
batch.sample_frame_index = [0]
|
|
batch.sample_num_frames = 1
|
|
return batch
|
|
|
|
assert batch.duration is not None and batch.fps is not None
|
|
batch.num_frames = math.ceil(batch.fps * batch.duration)
|
|
|
|
# Resample frame indices
|
|
frame_interval = batch.fps / self.train_fps
|
|
start_frame_idx = 0
|
|
frame_indices = np.arange(start_frame_idx, batch.num_frames,
|
|
frame_interval).astype(int)
|
|
|
|
# Temporal crop if too long
|
|
if len(frame_indices) > self.num_frames:
|
|
if temporal_sample_fn is not None:
|
|
begin_index, end_index = temporal_sample_fn(len(frame_indices))
|
|
frame_indices = frame_indices[begin_index:end_index]
|
|
else:
|
|
frame_indices = frame_indices[:self.num_frames]
|
|
|
|
batch.sample_frame_index = frame_indices.tolist()
|
|
batch.sample_num_frames = len(frame_indices)
|
|
|
|
return batch
|
|
|
|
|
|
class VideoTransformStage(DatasetStage):
|
|
"""Stage for video data transformation."""
|
|
|
|
def __init__(self, transform) -> None:
|
|
self.transform = transform
|
|
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""
|
|
Transform video data.
|
|
|
|
Args:
|
|
batch: Dataset batch with video information
|
|
|
|
Returns:
|
|
Batch with transformed video tensor
|
|
"""
|
|
if not batch.is_video:
|
|
return batch
|
|
|
|
assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
|
|
assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"
|
|
|
|
torchvision_video, _, metadata = torchvision.io.read_video(
|
|
batch.path, output_format="TCHW")
|
|
video = torchvision_video[batch.sample_frame_index]
|
|
if self.transform is not None:
|
|
video = self.transform(video)
|
|
video = rearrange(video, "t c h w -> c t h w")
|
|
video = video.to(torch.uint8)
|
|
|
|
video = video.float() / 127.5 - 1.0
|
|
batch.pixel_values = video
|
|
return batch
|
|
|
|
|
|
class ImageTransformStage(DatasetStage):
|
|
"""Stage for image data transformation."""
|
|
|
|
def __init__(self, transform, transform_topcrop) -> None:
|
|
self.transform = transform
|
|
self.transform_topcrop = transform_topcrop
|
|
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""
|
|
Transform image data.
|
|
|
|
Args:
|
|
batch: Dataset batch with image information
|
|
|
|
Returns:
|
|
Batch with transformed image tensor
|
|
"""
|
|
if not batch.is_image:
|
|
return batch
|
|
|
|
image = Image.open(batch.path).convert("RGB")
|
|
image = torch.from_numpy(np.array(image))
|
|
image = rearrange(image, "h w c -> c h w").unsqueeze(0)
|
|
|
|
if self.transform_topcrop is not None:
|
|
image = self.transform_topcrop(image)
|
|
elif self.transform is not None:
|
|
image = self.transform(image)
|
|
image = image.float() / 127.5 - 1.0
|
|
else:
|
|
image = image.float() / 127.5 - 1.0
|
|
|
|
image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W]
|
|
batch.pixel_values = image
|
|
return batch
|
|
|
|
|
|
class TextEncodingStage(DatasetStage):
|
|
"""Stage for text tokenization and encoding."""
|
|
|
|
def __init__(self,
|
|
tokenizer,
|
|
text_max_length: int,
|
|
cfg_rate: float = 0.0,
|
|
seed: int = 42):
|
|
self.tokenizer = tokenizer
|
|
self.text_max_length = text_max_length
|
|
self.cfg_rate = cfg_rate
|
|
# Create a seeded random generator for deterministic CFG
|
|
self.rng = random.Random(seed)
|
|
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""
|
|
Process text data.
|
|
|
|
Args:
|
|
batch: Dataset batch with caption information
|
|
|
|
Returns:
|
|
Batch with encoded text information
|
|
"""
|
|
text = batch.cap
|
|
if not isinstance(text, list):
|
|
text = [text]
|
|
text = [self.rng.choice(text)]
|
|
|
|
text = text[0] if self.rng.random() > self.cfg_rate else ""
|
|
text_tokens_and_mask = self.tokenizer(
|
|
text,
|
|
max_length=self.text_max_length,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_attention_mask=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
batch.text = text
|
|
batch.input_ids = text_tokens_and_mask["input_ids"]
|
|
batch.cond_mask = text_tokens_and_mask["attention_mask"]
|
|
return batch
|
|
|
|
|
|
class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
|
|
torch.distributed.checkpoint.stateful.Stateful):
|
|
"""
|
|
Merged dataset for video and caption data with stage-based processing.
|
|
Assumes that data_merge_path is a txt file with the following format:
|
|
<folder_path>,<json_file_path>
|
|
|
|
The folder should contain videos.
|
|
|
|
The json file should be a list of dictionaries with the following format:
|
|
[
|
|
{
|
|
"path": "1gGQy4nxyUo-Scene-016.mp4",
|
|
"resolution": {
|
|
"width": 1920,
|
|
"height": 1080
|
|
},
|
|
"size": 2439112,
|
|
"fps": 25.0,
|
|
"duration": 6.88,
|
|
"num_frames": 172,
|
|
"cap": [
|
|
"A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
|
|
]
|
|
},
|
|
...
|
|
]
|
|
|
|
This dataset processes video and image data through a series of stages:
|
|
- Data validation
|
|
- Resolution filtering
|
|
- Frame sampling
|
|
- Transformation
|
|
- Text encoding
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_merge_path: str,
|
|
args,
|
|
transform,
|
|
temporal_sample,
|
|
transform_topcrop,
|
|
start_idx: int = 0,
|
|
seed: int = 42):
|
|
self.data_merge_path = data_merge_path
|
|
self.start_idx = start_idx
|
|
self.args = args
|
|
self.temporal_sample = temporal_sample
|
|
self.seed = seed
|
|
|
|
# Initialize tokenizer
|
|
tokenizer_path = os.path.join(args.model_path, "tokenizer")
|
|
tokenizer = None
|
|
if os.path.exists(tokenizer_path):
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
tokenizer_path, cache_dir=args.cache_dir)
|
|
except (ValueError, OSError):
|
|
pass
|
|
|
|
# Initialize processing stages
|
|
self._init_stages(args, transform, transform_topcrop, tokenizer)
|
|
|
|
# Process metadata
|
|
self.processed_batches = self._process_metadata()
|
|
|
|
def _init_stages(self, args, transform, transform_topcrop,
|
|
tokenizer) -> None:
|
|
"""Initialize all processing stages."""
|
|
self.validation_stage = DataValidationStage()
|
|
self.frame_sampling_stage = FrameSamplingStage(
|
|
num_frames=args.num_frames,
|
|
train_fps=args.train_fps,
|
|
speed_factor=args.speed_factor,
|
|
video_length_tolerance_range=args.video_length_tolerance_range,
|
|
drop_short_ratio=args.drop_short_ratio,
|
|
seed=self.seed)
|
|
self.video_transform_stage = VideoTransformStage(transform)
|
|
self.image_transform_stage = ImageTransformStage(
|
|
transform, transform_topcrop)
|
|
if tokenizer is not None:
|
|
self.text_encoding_stage = TextEncodingStage(
|
|
tokenizer=tokenizer,
|
|
text_max_length=args.text_max_length,
|
|
cfg_rate=args.training_cfg_rate,
|
|
seed=self.seed)
|
|
else:
|
|
self.text_encoding_stage = None
|
|
|
|
def _load_raw_data(self) -> list[dict]:
|
|
"""Load raw data from JSON files."""
|
|
# Read folder-annotation pairs
|
|
with open(self.data_merge_path) as f:
|
|
folder_anno_pairs = [
|
|
line.strip().split(",") for line in f if line.strip()
|
|
]
|
|
assert len(
|
|
folder_anno_pairs) == 1, "Only support one folder-annotation pair"
|
|
assert len(folder_anno_pairs[0]
|
|
) == 2, "Folder-annotation pair should have two elements"
|
|
folder, annotation_file = folder_anno_pairs[0]
|
|
|
|
data_items: list[dict] = []
|
|
with open(annotation_file) as f:
|
|
data_items = json.load(f)
|
|
|
|
# Update paths with folder prefix
|
|
for item in data_items:
|
|
item["path"] = opj(folder, item["path"])
|
|
if "action_path" in item and item["action_path"]:
|
|
item["action_path"] = opj(folder, item["action_path"])
|
|
|
|
return data_items
|
|
|
|
def _process_metadata(self) -> list[PreprocessBatch]:
|
|
"""Process the raw metadata through all filtering stages."""
|
|
raw_data = self._load_raw_data()
|
|
processed_batches = []
|
|
|
|
# Initialize counters
|
|
filter_counts = {
|
|
"validation_failed": 0,
|
|
"frame_sampling_failed": 0
|
|
}
|
|
sample_num_frames: list[int] = []
|
|
|
|
for item in raw_data:
|
|
batch = PreprocessBatch(path=item["path"],
|
|
cap=item["cap"],
|
|
resolution=item.get("resolution"),
|
|
fps=item.get("fps"),
|
|
duration=item.get("duration"),
|
|
action_path=item.get("action_path"))
|
|
|
|
# Apply filtering stages
|
|
if not self._apply_filter_stages(batch, filter_counts):
|
|
continue
|
|
|
|
# Apply frame sampling processing
|
|
batch = self.frame_sampling_stage.process(
|
|
batch, temporal_sample_fn=self.temporal_sample)
|
|
|
|
processed_batches.append(batch)
|
|
assert batch.sample_num_frames is not None
|
|
sample_num_frames.append(batch.sample_num_frames)
|
|
|
|
self._log_filtering_stats(filter_counts, sample_num_frames,
|
|
len(raw_data), len(processed_batches))
|
|
return processed_batches
|
|
|
|
def _apply_filter_stages(self, batch: PreprocessBatch,
|
|
filter_counts: dict[str, int]) -> bool:
|
|
"""Apply all filter stages and update counters. Returns True if batch should be kept."""
|
|
if not self.validation_stage.should_keep(batch):
|
|
filter_counts["validation_failed"] += 1
|
|
return False
|
|
|
|
if not self.frame_sampling_stage.should_keep(batch):
|
|
filter_counts["frame_sampling_failed"] += 1
|
|
return False
|
|
|
|
return True
|
|
|
|
def _log_filtering_stats(self, filter_counts: dict[str, int],
|
|
sample_num_frames: list[int], before_count: int,
|
|
after_count: int):
|
|
"""Log filtering statistics."""
|
|
logger.info(
|
|
"validation_failed: %d, frame_sampling_failed: %d, "
|
|
"Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
|
|
filter_counts['validation_failed'],
|
|
filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
|
|
before_count, after_count)
|
|
|
|
def __iter__(self):
|
|
"""Iterate through processed data items."""
|
|
for idx in range(len(self.processed_batches)):
|
|
yield self._get_item(idx)
|
|
|
|
def __len__(self):
|
|
return len(self.processed_batches)
|
|
|
|
def _get_item(self, idx: int) -> dict:
|
|
"""Get a single processed data item."""
|
|
batch = self.processed_batches[idx]
|
|
|
|
# Apply transformation stages
|
|
batch = self.video_transform_stage.process(batch)
|
|
batch = self.image_transform_stage.process(batch)
|
|
if self.text_encoding_stage is not None:
|
|
batch = self.text_encoding_stage.process(batch)
|
|
|
|
# Build result dictionary
|
|
result = {
|
|
"pixel_values": batch.pixel_values,
|
|
"path": batch.path,
|
|
}
|
|
|
|
if batch.text is not None:
|
|
result["text"] = batch.text
|
|
result["input_ids"] = batch.input_ids
|
|
result["cond_mask"] = batch.cond_mask
|
|
|
|
# Add video-specific fields
|
|
if batch.is_video:
|
|
result.update({"fps": batch.fps, "duration": batch.duration})
|
|
|
|
# Add action_path
|
|
if batch.action_path:
|
|
result["action_path"] = batch.action_path
|
|
|
|
return result
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
"""Return state dict for checkpointing."""
|
|
return {"processed_batches": self.processed_batches}
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
"""Load state dict from checkpoint."""
|
|
self.processed_batches = state_dict["processed_batches"]
|
|
|
|
|
|
class TextDataset(torch.utils.data.IterableDataset,
|
|
torch.distributed.checkpoint.stateful.Stateful):
|
|
"""
|
|
Text-only dataset for processing prompts from a simple text file.
|
|
|
|
Assumes that data_merge_path is a text file with one prompt per line:
|
|
A cat playing with a ball
|
|
A dog running in the park
|
|
A person cooking dinner
|
|
...
|
|
|
|
This dataset processes text data through text encoding stages only.
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_merge_path: str,
|
|
args,
|
|
start_idx: int = 0,
|
|
seed: int = 42):
|
|
self.data_merge_path = data_merge_path
|
|
self.start_idx = start_idx
|
|
self.args = args
|
|
self.seed = seed
|
|
|
|
# Initialize tokenizer
|
|
tokenizer_path = os.path.join(args.model_path, "tokenizer")
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
|
|
cache_dir=args.cache_dir)
|
|
|
|
# Initialize text encoding stage
|
|
self.text_encoding_stage = TextEncodingStage(
|
|
tokenizer=tokenizer,
|
|
text_max_length=args.text_max_length,
|
|
cfg_rate=getattr(args, 'training_cfg_rate', 0.0),
|
|
seed=self.seed)
|
|
|
|
# Process text data
|
|
self.processed_batches = self._process_text_data()
|
|
|
|
def _load_text_data(self) -> list[str]:
|
|
"""Load text prompts from file."""
|
|
prompts = []
|
|
with open(self.data_merge_path, 'r', encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line: # Skip empty lines
|
|
prompts.append(line)
|
|
|
|
logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}")
|
|
return prompts
|
|
|
|
def _process_text_data(self) -> list[PreprocessBatch]:
|
|
"""Process the text prompts through text encoding stage."""
|
|
raw_prompts = self._load_text_data()
|
|
processed_batches = []
|
|
|
|
for idx, prompt in enumerate(raw_prompts):
|
|
# Create a text-only batch with dummy path
|
|
batch = PreprocessBatch(
|
|
path=f"text_prompt_{idx}",
|
|
cap=[prompt], # TextEncodingStage expects a list
|
|
resolution=None,
|
|
fps=None,
|
|
duration=None,
|
|
num_frames=0,
|
|
sample_frame_index=None,
|
|
sample_num_frames=0
|
|
)
|
|
|
|
processed_batches.append(batch)
|
|
|
|
logger.info(f"Processed {len(processed_batches)} text batches")
|
|
return processed_batches
|
|
|
|
def __iter__(self):
|
|
"""Iterator for the dataset."""
|
|
# Set up distributed sampling if needed
|
|
if torch.distributed.is_available() and torch.distributed.is_initialized():
|
|
rank = torch.distributed.get_rank()
|
|
world_size = torch.distributed.get_world_size()
|
|
else:
|
|
rank = 0
|
|
world_size = 1
|
|
|
|
# Calculate chunk for this rank
|
|
total_items = len(self.processed_batches)
|
|
items_per_rank = math.ceil(total_items / world_size)
|
|
start_idx = rank * items_per_rank + self.start_idx
|
|
end_idx = min(start_idx + items_per_rank, total_items)
|
|
|
|
# Yield items for this rank
|
|
for idx in range(start_idx, end_idx):
|
|
if idx < len(self.processed_batches):
|
|
yield self._get_item(idx)
|
|
|
|
def _get_item(self, idx: int) -> dict:
|
|
"""Get a single processed text item."""
|
|
batch = self.processed_batches[idx]
|
|
|
|
# Apply text encoding stage
|
|
batch = self.text_encoding_stage.process(batch)
|
|
|
|
# Build result dictionary for text-only processing with required schema fields
|
|
result = {
|
|
"text": batch.text,
|
|
"input_ids": batch.input_ids,
|
|
"cond_mask": batch.cond_mask,
|
|
"path": batch.path,
|
|
# Required schema fields for ODE trajectory processing
|
|
"id": f"text_{idx}",
|
|
"file_name": batch.path,
|
|
"caption": batch.text,
|
|
"media_type": "text",
|
|
"width": 1,
|
|
"height": 1,
|
|
"num_frames": 0,
|
|
"duration_sec": 0.0,
|
|
"fps": 0.0,
|
|
}
|
|
|
|
return result
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
"""Return state dict for checkpointing."""
|
|
return {"processed_batches": self.processed_batches}
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
"""Load state dict from checkpoint."""
|
|
self.processed_batches = state_dict["processed_batches"]
|