caption-flow 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

caption_flow/utils/dataset_metadata_cache.py

@@ -0,0 +1,67 @@
+"""Dataset metadata caching for efficient HuggingFace dataset handling."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional, List
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetMetadataCache:
+    """Caches dataset metadata to avoid repeated full iterations."""
+
+    def __init__(self, cache_dir: Path):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.cache_file = self.cache_dir / "dataset_metadata.json"
+        self.metadata: Dict[str, Any] = {}
+        self._load_cache()
+
+    def _load_cache(self):
+        """Load cached metadata from disk."""
+        if self.cache_file.exists():
+            try:
+                with open(self.cache_file, "r") as f:
+                    self.metadata = json.load(f)
+                logger.info(f"Loaded dataset metadata cache with {len(self.metadata)} datasets")
+            except Exception as e:
+                logger.error(f"Failed to load metadata cache: {e}")
+                self.metadata = {}
+
+    def _save_cache(self):
+        """Save metadata cache to disk."""
+        try:
+            with open(self.cache_file, "w") as f:
+                json.dump(self.metadata, f, indent=2)
+            logger.debug("Saved dataset metadata cache")
+        except Exception as e:
+            logger.error(f"Failed to save metadata cache: {e}")
+
+    def get_dataset_key(self, dataset_path: str, split: str) -> str:
+        """Generate a unique key for a dataset+split combination."""
+        return f"{dataset_path}:{split}"
+
+    def get_metadata(self, dataset_path: str, split: str) -> Optional[Dict[str, Any]]:
+        """Get cached metadata for a dataset."""
+        key = self.get_dataset_key(dataset_path, split)
+        return self.metadata.get(key)
+
+    def set_metadata(self, dataset_path: str, split: str, metadata: Dict[str, Any]):
+        """Cache metadata for a dataset."""
+        key = self.get_dataset_key(dataset_path, split)
+        metadata["cached_at"] = datetime.utcnow().isoformat()
+        metadata["dataset_path"] = dataset_path
+        metadata["split"] = split
+        self.metadata[key] = metadata
+        self._save_cache()
+        logger.info(f"Cached metadata for {key}: {metadata.get('total_items', 0)} items")
+
+    def invalidate(self, dataset_path: str, split: str):
+        """Remove cached metadata for a dataset."""
+        key = self.get_dataset_key(dataset_path, split)
+        if key in self.metadata:
+            del self.metadata[key]
+            self._save_cache()
+            logger.info(f"Invalidated metadata cache for {key}")

caption_flow/utils/image_processor.py

@@ -112,10 +112,7 @@ class ImageProcessor:
                 return None
 
         except Exception as e:
-            logger.error(f"Error processing image data: {e}")
-            import traceback
-
-            logger.error(traceback.format_exc())
+            logger.error(f"Error processing image data: {e}", exc_info=True)
             return None
 
     @staticmethod

caption_flow/utils/shard_processor.py

@@ -7,7 +7,6 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Generator, Tuple, Optional, Dict, Any
 from dataclasses import dataclass
-from datasets import load_dataset
 from .image_processor import ImageProcessor
 from threading import Event
 import shlex
@@ -24,213 +23,22 @@ class ShardProcessor(ABC):
     """Abstract base for processing dataset shards."""
 
     @abstractmethod
-    def iterate_chunk(
+    def iterate_chunk_with_metadata(
         self,
         chunk,
         dataset_loader: Optional[DatasetLoader],
         should_stop: Event,
         connected: Event,
-    ) -> Generator[Tuple[str, str, bytes], None, None]:
+    ) -> Generator[Tuple[str, str, bytes, Dict[str, Any]], None, None]:
         """
-        Iterate through items in a chunk.
+        Iterate through items in a chunk with metadata.
 
         Yields:
-            Tuple of (key, url, image_data)
+            Tuple of (key, url, image_data, metadata)
         """
         pass
 
 
-class HFDatasetShardProcessor(ShardProcessor):
-    """Processor for HuggingFace virtual dataset shards."""
-
-    def iterate_chunk(
-        self,
-        chunk,
-        dataset_loader: Optional[DatasetLoader],
-        should_stop: Event,
-        connected: Event,
-    ) -> Generator[Tuple[str, str, bytes], None, None]:
-        """Process HuggingFace virtual shard chunk."""
-        if not dataset_loader:
-            logger.error("No dataset loader configured for HuggingFace dataset shard")
-            return
-
-        # Get unprocessed ranges
-        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
-
-        logger.info(
-            f"Processing HF dataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
-        )
-
-        items_processed = 0
-        current_idx = 0
-
-        # Construct proper virtual shard URL
-        parts = chunk.shard_url.split("_chunk_")
-        if len(parts) == 2:
-            base_path = parts[0]
-            virtual_shard_url = f"{base_path}:chunk:{chunk.start_index}"
-        else:
-            virtual_shard_url = chunk.shard_url
-
-        logger.debug(f"Using virtual shard URL: {virtual_shard_url}")
-
-        # Iterate through the virtual shard
-        for key, url, image_data in dataset_loader.iterate_shard(virtual_shard_url):
-            # Check if we should stop
-            if should_stop.is_set() or not connected.is_set():
-                logger.info(f"Stopping chunk processing early due to disconnect")
-                break
-
-            # Check if current index is in any unprocessed range
-            in_range = any(start <= current_idx <= end for start, end in unprocessed_ranges)
-
-            if not in_range:
-                current_idx += 1
-                continue  # Skip already processed items
-
-            # Check if we've processed enough for this chunk
-            if current_idx >= chunk.chunk_size:
-                break
-
-            items_processed += 1
-            current_idx += 1
-            yield key, url, image_data
-
-        logger.info(
-            f"HF dataset chunk {chunk.chunk_id}: yielded {items_processed} items "
-            f"from ranges {unprocessed_ranges}"
-        )
-
-    def iterate_chunk_with_metadata(
-        self,
-        chunk,
-        dataset_loader: Optional[DatasetLoader],
-        should_stop: Event,
-        connected: Event,
-    ) -> Generator[Tuple[str, str, bytes, Dict[str, Any]], None, None]:
-        """
-        Process HuggingFace virtual shard chunk with metadata, range by range.
-        """
-        if not dataset_loader:
-            logger.error("No dataset loader configured for HuggingFace dataset shard")
-            return
-
-        # Get unprocessed ranges
-        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
-
-        logger.info(
-            f"Processing HF dataset chunk {chunk.chunk_id} with {len(unprocessed_ranges)} ranges"
-        )
-
-        items_yielded = 0
-
-        # Process each range independently with its own iterator
-        for range_start, range_end in unprocessed_ranges:
-            if should_stop.is_set() or not connected.is_set():
-                logger.info(f"Stopping chunk processing early due to disconnect")
-                break
-
-            # Calculate absolute indices for this range
-            abs_start = chunk.start_index + range_start
-            abs_end = chunk.start_index + range_end
-            range_size = range_end - range_start + 1
-
-            logger.debug(
-                f"Processing range [{range_start}, {range_end}] "
-                f"(absolute: [{abs_start}, {abs_end}])"
-            )
-
-            try:
-                # Create a fresh dataset iterator for this range
-                dataset = load_dataset(
-                    dataset_loader.dataset_path,
-                    split=dataset_loader.split,
-                    streaming=True,
-                    token=dataset_loader.token,
-                )
-
-                # Use state_dict if available for efficient positioning
-                if hasattr(dataset, "load_state_dict") and hasattr(dataset, "state_dict"):
-                    try:
-                        state = dataset.state_dict()
-                        # Modify state to jump to abs_start
-                        if "num_examples_since_previous_state" in state:
-                            state["num_examples_since_previous_state"] = abs_start
-                        if "examples_iterable" in state and isinstance(
-                            state["examples_iterable"], dict
-                        ):
-                            if "shard_example_idx" in state["examples_iterable"]:
-                                state["examples_iterable"]["shard_example_idx"] = abs_start
-                        dataset.load_state_dict(state)
-                        logger.debug(f"Positioned dataset at index {abs_start} using state_dict")
-                    except Exception as e:
-                        logger.debug(f"Could not use state_dict, falling back to skip: {e}")
-                        dataset = dataset.skip(abs_start)
-                else:
-                    # Fall back to skip
-                    dataset = dataset.skip(abs_start)
-
-                # Process items in this range
-                range_items = 0
-                for item in dataset:
-                    if range_items >= range_size:
-                        break
-
-                    if should_stop.is_set() or not connected.is_set():
-                        break
-
-                    # Generate key for this item
-                    current_abs_idx = abs_start + range_items
-                    key = f"{dataset_loader.dataset_path.replace('/', '_')}_{current_abs_idx:08d}"
-
-                    try:
-                        if dataset_loader.image_column in item:
-                            img_data = item[dataset_loader.image_column]
-                            image_bytes = ImageProcessor.process_image_data(img_data)
-
-                            if image_bytes:
-                                # Extract metadata
-                                metadata = {
-                                    k: v
-                                    for k, v in item.items()
-                                    if k != dataset_loader.image_column
-                                }
-                                # Add chunk-relative index to metadata
-                                metadata["_chunk_relative_index"] = range_start + range_items
-
-                                url = f"hf://{dataset_loader.dataset_path}#{current_abs_idx}"
-
-                                items_yielded += 1
-                                range_items += 1
-
-                                yield key, url, image_bytes, metadata
-                            else:
-                                logger.warning(
-                                    f"Failed to process image at index {current_abs_idx}"
-                                )
-                                range_items += 1
-                        else:
-                            logger.warning(
-                                f"No image column '{dataset_loader.image_column}' at index {current_abs_idx}"
-                            )
-                            range_items += 1
-
-                    except Exception as e:
-                        logger.error(f"Error processing item at index {current_abs_idx}: {e}")
-                        range_items += 1
-                        continue
-
-            except Exception as e:
-                logger.error(f"Error processing range [{range_start}, {range_end}]: {e}")
-                continue
-
-        logger.info(
-            f"HF dataset chunk {chunk.chunk_id}: yielded {items_yielded} items "
-            f"from {len(unprocessed_ranges)} ranges"
-        )
-
-
 class WebDatasetShardProcessor(ShardProcessor):
     """Processor for WebDataset tar shards with range support."""
 
@@ -238,74 +46,6 @@ class WebDatasetShardProcessor(ShardProcessor):
         self.hf_token = hf_token
         self.dataset_type = dataset_type
 
-    def iterate_chunk(
-        self,
-        chunk,
-        dataset_loader: Optional[DatasetLoader],
-        should_stop: Event,
-        connected: Event,
-    ) -> Generator[Tuple[str, str, bytes], None, None]:
-        """Process WebDataset shard chunk with unprocessed ranges."""
-        # Get unprocessed ranges
-        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
-
-        logger.info(
-            f"Processing WebDataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
-        )
-
-        # Create WebDataset pipeline
-        if self.dataset_type == "huggingface" and not chunk.shard_url.startswith("hf_dataset:"):
-            # Use curl with auth for HuggingFace WebDataset
-            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.hf_token)}' {shlex.quote(chunk.shard_url)} || true"
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(url_cmd),
-                wds.tarfile_to_samples(),
-                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
-            )
-        else:
-            # Local file
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(chunk.shard_url),
-                wds.tarfile_to_samples(),
-                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
-            )
-
-        # Process items
-        current_idx = 0
-        items_yielded = 0
-
-        for key, image_data in ds:
-            # Check if we should stop
-            if should_stop.is_set() or not connected.is_set():
-                logger.info(f"Stopping WebDataset chunk processing early due to disconnect")
-                break
-
-            # Calculate relative index within chunk
-            relative_idx = current_idx - chunk.start_index
-
-            # Skip items before chunk start
-            if current_idx < chunk.start_index:
-                current_idx += 1
-                continue
-
-            # Stop if beyond chunk
-            if relative_idx >= chunk.chunk_size:
-                break
-
-            # Check if current index is in any unprocessed range
-            in_range = any(start <= relative_idx <= end for start, end in unprocessed_ranges)
-
-            if in_range:
-                items_yielded += 1
-                yield key, chunk.shard_url, image_data
-
-            current_idx += 1
-
-        logger.info(
-            f"WebDataset chunk {chunk.chunk_id}: yielded {items_yielded} items "
-            f"from ranges {unprocessed_ranges}"
-        )
-
     def iterate_chunk_with_metadata(
         self,
         chunk,
@@ -322,7 +62,7 @@ class WebDatasetShardProcessor(ShardProcessor):
         )
 
         # Create WebDataset pipeline
-        if self.dataset_type == "huggingface" and not chunk.shard_url.startswith("hf_dataset:"):
+        if self.dataset_type == "huggingface":
             # Use curl with auth for HuggingFace WebDataset
             url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.hf_token)}' {shlex.quote(chunk.shard_url)} || true"
             ds = wds.DataPipeline(

caption_flow/utils/shard_tracker.py

@@ -61,11 +61,7 @@ class ShardTracker(CheckpointTracker):
         """Get list of shards that still need processing."""
         remaining = []
         for s in all_shards:
-            # Extract shard name properly for both regular and virtual shards
-            if s.startswith("hf_dataset:"):
-                shard_name = s  # Use full virtual shard ID
-            else:
-                shard_name = Path(s).stem
+            shard_name = Path(s).stem
 
             if shard_name not in self.completed_shards:
                 remaining.append(s)