Video support for datatrove #271
Conversation
There was a problem hiding this comment.
Not sure how much detail you wanted in the feedback, some stuff are nits but I included them just in case.
Main issue imo is no sharding on the reader which would be a problem if you want to scale up/use multiple cpus. It's easily fixible, basically just pass rank and world_size to your find_triplets function and change the return to return triplets[rank::world_size]. This should be consistent across workers as list_files is deterministic (it sorts the file list) so there should be no overlap between different workers
| """Filter that uses ffmpeg to detect if a video is static (frozen).""" | ||
|
|
||
| name = "🧊 Video-Frozen-filter" | ||
| _requires_dependencies = ["ffmpeg"] |
There was a problem hiding this comment.
Should probably be this to show users the correct package name to install
| _requires_dependencies = ["ffmpeg"] | |
| _requires_dependencies = [("ffmpeg", "ffmpeg-python")] |
There was a problem hiding this comment.
You forgot to accept the actual change here I think
| self.data_folder = get_datafolder(data_folder) | ||
| self.paths_file = paths_file |
There was a problem hiding this comment.
you don't need to save self.data_folder and self.paths_file as these are already saved by the base class
| """Overrides the base run method to handle triplet statistics correctly.""" | ||
| triplet_count = 0 | ||
|
|
||
| for triplet in self.find_triplets(): |
There was a problem hiding this comment.
You're currently not sharding, which will make it hard to parallelize/easily scale up to many cpus
| with self.data_folder.open(self.paths_file, "r") as f: | ||
| paths = [line.strip() for line in f] | ||
| else: | ||
| paths = self.data_folder.list_files(recursive=self.recursive, glob_pattern="*") |
There was a problem hiding this comment.
if you don't want to use the glob_pattern you can just pass None (or nothing as it is the default), but then you probably want to remove it from the __init__ args
|
|
||
| for path in paths: | ||
| base_name, ext = os.path.splitext(path) | ||
| if ext in video_extensions: |
There was a problem hiding this comment.
we could maybe extract the triple finding thing away to the datafolder class later on (you'd pass a tuple like ((".mp4", ".avi", ".mkv", ".mov"), ".json", ".vtt")), but it's probably ok for now
| def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline: | ||
| """Overrides the base run method to handle triplet statistics correctly.""" | ||
| triplet_count = 0 | ||
|
|
There was a problem hiding this comment.
This allows stacking multiple readers together (when you want to load data from multiple sources for example)
| if data: | |
| yield from data |
|
|
||
| return document | ||
|
|
||
| def ensure_local_copy(self, video_url: str) -> str: |
There was a problem hiding this comment.
good job using all of datatrove useful io features here :)
| """Filter that uses ffmpeg to detect if a video is static (frozen).""" | ||
|
|
||
| name = "🧊 Video-Frozen-filter" | ||
| _requires_dependencies = ["ffmpeg"] |
There was a problem hiding this comment.
You forgot to accept the actual change here I think
WIP