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

caption_flow/utils/dataset_loader.py

@@ -1,222 +0,0 @@
-"""Dataset loading utilities for WebDataset and HuggingFace."""
-
-import asyncio
-import shlex
-import logging
-from pathlib import Path
-from typing import List, Dict, Any, Generator, Optional, Tuple
-import json
-
-import webdataset as wds
-from huggingface_hub import HfFileSystem, get_token, hf_hub_url
-
-logger = logging.getLogger(__name__)
-
-
-class DatasetLoader:
-    """Handles loading datasets from various sources."""
-
-    def __init__(
-        self,
-        dataset_path: str,
-        dataset_type: str = "huggingface",
-        split: str = "train",
-        image_column: str = "image",
-        cache_dir: Optional[Path] = None,
-    ):
-        """
-        Initialize dataset loader.
-
-        Args:
-            dataset_path: Path to dataset (HF repo, local dir, etc.)
-            dataset_type: Type of dataset ("huggingface", "webdataset", "local")
-            split: Split to use for HuggingFace datasets (default: "train")
-            image_column: Column name containing image data or URLs (default: "image")
-        """
-        self.dataset_path = dataset_path
-        self.dataset_type = dataset_type
-        self.split = split
-        self.image_column = image_column
-        self.token = get_token()
-        self.dataset_format = None  # Will be detected: "webdataset" or "huggingface_datasets"
-
-        if not self.token and dataset_type == "huggingface":
-            logger.warning("No HuggingFace token found; run `huggingface-cli login`")
-
-        # Detect the actual format if it's a HuggingFace dataset
-        if dataset_type == "huggingface":
-            self.dataset_format = self._detect_dataset_format()
-            logger.info(f"Detected dataset format: {self.dataset_format}")
-
-    def _detect_dataset_format(self) -> str:
-        """Detect whether it's WebDataset or HuggingFace datasets format."""
-        fs = HfFileSystem(token=self.token)
-
-        # Check for .tar files (WebDataset)
-        tar_files = list(fs.glob(f"hf://datasets/{self.dataset_path}/**/*.tar"))
-        if tar_files:
-            return "webdataset"
-
-        # Check for .parquet files (Huggingface Arrow DB)
-        parquet_files = list(fs.glob(f"hf://datasets/{self.dataset_path}/**/*.parquet"))
-        if parquet_files:
-            return "huggingface_datasets"
-
-        raise AssertionError(f"Could not detect dataset format for {self.dataset_path}")
-
-    def get_shard_list(self) -> List[str]:
-        """Get list of all shards in the dataset."""
-        if self.dataset_type == "huggingface":
-            if self.dataset_format == "webdataset":
-                return self._get_hf_webdataset_shards()
-            else:
-                logger.error(f"Unknown dataset format: {self.dataset_format}")
-                return []
-        elif self.dataset_type == "local":
-            return self._get_local_shards()
-        else:
-            raise ValueError(f"Unknown dataset type: {self.dataset_type}")
-
-    def _get_hf_webdataset_shards(self) -> List[str]:
-        """Get shard URLs from HuggingFace WebDataset."""
-        logger.info(f"Getting WebDataset shard list from HuggingFace: {self.dataset_path}")
-
-        fs = HfFileSystem(token=self.token)
-        files = [fs.resolve_path(p) for p in fs.glob(f"hf://datasets/{self.dataset_path}/**/*.tar")]
-
-        urls = [hf_hub_url(f.repo_id, f.path_in_repo, repo_type="dataset") for f in files]
-
-        logger.info(f"Found {len(urls)} WebDataset shards")
-        return sorted(urls)
-
-    def _get_local_shards(self) -> List[str]:
-        """Get shard files from local directory."""
-        path = Path(self.dataset_path)
-        if not path.exists():
-            raise ValueError(f"Local dataset path does not exist: {path}")
-
-        shards = list(path.glob("*.tar"))
-        logger.info(f"Found {len(shards)} local shards")
-        return [str(s) for s in sorted(shards)]
-
-    def load_shard(self, shard_url: str, processed_keys: Optional[set] = None) -> wds.DataPipeline:
-        """
-        Load a single shard as a WebDataset pipeline.
-
-        Args:
-            shard_url: URL or path to the shard
-            processed_keys: Set of already processed keys to skip
-        """
-        if processed_keys is None:
-            processed_keys = set()
-
-        if self.dataset_type == "huggingface" and self.dataset_format == "webdataset":
-            # Use curl with auth token for HuggingFace
-            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.token)}' {shlex.quote(shard_url)} || true"
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(url_cmd),
-                wds.tarfile_to_samples(),
-                wds.to_tuple("__key__", "__url__", "jpg;png;jpeg;webp;jxl"),
-                wds.select(lambda x: x[0] not in processed_keys),
-            )
-        else:
-            # Local file access
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(shard_url),
-                wds.tarfile_to_samples(),
-                wds.to_tuple("__key__", "__url__", "jpg;png;jpeg;webp;jxl"),
-                wds.select(lambda x: x[0] not in processed_keys),
-            )
-
-        return ds
-
-    def iterate_shard(
-        self,
-        shard_url: str,
-        processed_keys: Optional[set] = None,
-        unprocessed_ranges: Optional[List[Tuple[int, int]]] = None,
-    ) -> Generator[Dict[str, Any], None, None]:
-        """
-        Iterate over items in a shard, returning full sample dictionaries.
-
-        Args:
-            shard_url: URL or identifier of the shard
-            processed_keys: Set of already processed keys to skip
-            unprocessed_ranges: Specific ranges to process (for range-based processing)
-
-        Yields:
-            Dictionary containing the full WebDataset sample
-        """
-        if processed_keys is None:
-            processed_keys = set()
-
-        if self.dataset_type == "huggingface" and self.dataset_format == "webdataset":
-            # Use curl with auth token for HuggingFace
-            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.token)}' {shlex.quote(shard_url)} || true"
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(url_cmd),
-                wds.tarfile_to_samples(),
-                wds.select(lambda x: x.get("__key__", "") not in processed_keys),
-            )
-        else:
-            # Local file access
-            ds = wds.DataPipeline(
-                wds.SimpleShardList(shard_url),
-                wds.tarfile_to_samples(),
-                wds.select(lambda x: x.get("__key__", "") not in processed_keys),
-            )
-
-        # Return full samples as dictionaries
-        for sample in ds:
-            # Ensure it's a dict and has required fields
-            if isinstance(sample, dict) and "__key__" in sample:
-                yield sample
-
-    def count_shard_items(self, shard_url: str, processed_keys: Optional[set] = None) -> int:
-        """Count items in a shard (can be slow for large shards)."""
-        count = 0
-        try:
-            for _ in self.iterate_shard(shard_url, processed_keys):
-                count += 1
-        except Exception as e:
-            logger.error(f"Error counting shard {shard_url}: {e}")
-        return count
-
-    def get_dataset_info(self) -> Dict[str, Any]:
-        """Get information about the dataset."""
-        info = {
-            "dataset_path": self.dataset_path,
-            "dataset_type": self.dataset_type,
-            "dataset_format": self.dataset_format,
-        }
-
-        if self.dataset_format == "huggingface_datasets":
-            # Include cached metadata if available
-            if hasattr(self, "_hf_metadata"):
-                info.update(self._hf_metadata)
-            else:
-
-                try:
-                    # Try to get more info about the dataset
-                    dataset_info = load_dataset(
-                        self.dataset_path, split=self.split, streaming=True, token=self.token
-                    )
-                    # Get features info
-                    if hasattr(dataset_info, "features"):
-                        info["features"] = str(dataset_info.features)
-
-                    # Try to get total size (might not work for all datasets)
-                    try:
-                        # This might be expensive for large datasets
-                        total_examples = len(
-                            load_dataset(self.dataset_path, split=self.split, token=self.token)
-                        )
-                        info["total_examples"] = total_examples
-                        self._hf_total_items = total_examples
-                    except:
-                        info["total_examples"] = "unknown"
-
-                except Exception as e:
-                    logger.error(f"Error getting dataset info: {e}")
-
-        return info

caption_flow/utils/dataset_metadata_cache.py

@@ -1,67 +0,0 @@
-"""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/job_queue.py

@@ -1,41 +0,0 @@
-"""Job queue management."""
-
-import asyncio
-from typing import Optional
-from collections import deque
-
-from ..models import Job
-
-
-class JobQueue:
-    """Priority job queue with backpressure."""
-
-    def __init__(self):
-        self.queue = deque()
-        self.processing = set()
-        self.lock = asyncio.Lock()
-
-    async def add(self, job: Job):
-        """Add job to queue."""
-        async with self.lock:
-            self.queue.append(job)
-
-    async def get_next(self) -> Optional[Job]:
-        """Get next available job."""
-        async with self.lock:
-            if self.queue:
-                job = self.queue.popleft()
-                self.processing.add(job.job_id)
-                return job
-        return None
-
-    async def complete(self, job_id: str):
-        """Mark job as complete."""
-        async with self.lock:
-            self.processing.discard(job_id)
-
-    async def requeue(self, job: Job):
-        """Requeue a job (for failures)."""
-        async with self.lock:
-            self.processing.discard(job.job_id)
-            self.queue.appendleft(job)  # Priority requeue

caption_flow/utils/shard_processor.py

@@ -1,119 +0,0 @@
-"""Shard processing abstraction for different dataset types."""
-
-import io
-import logging
-import time
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Generator, Tuple, Optional, Dict, Any
-from dataclasses import dataclass
-from .image_processor import ImageProcessor
-from threading import Event
-import shlex
-
-import webdataset as wds
-from PIL import Image
-
-from .dataset_loader import DatasetLoader
-
-logger = logging.getLogger(__name__)
-
-
-class ShardProcessor(ABC):
-    """Abstract base for processing dataset shards."""
-
-    @abstractmethod
-    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]:
-        """
-        Iterate through items in a chunk with metadata.
-
-        Yields:
-            Tuple of (key, url, image_data, metadata)
-        """
-        pass
-
-
-class WebDatasetShardProcessor(ShardProcessor):
-    """Processor for WebDataset tar shards with range support."""
-
-    def __init__(self, hf_token: Optional[str] = None, dataset_type: str = "local"):
-        self.hf_token = hf_token
-        self.dataset_type = dataset_type
-
-    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 WebDataset shard chunk with metadata and range support."""
-        # 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":
-            # 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
-        absolute_idx = 0  # Absolute index in the shard
-        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
-
-            # Skip items before chunk start
-            if absolute_idx < chunk.start_index:
-                absolute_idx += 1
-                continue
-
-            # Calculate relative index within chunk
-            relative_idx = absolute_idx - chunk.start_index
-
-            # 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:
-                # Create metadata with the relative index
-                metadata = {
-                    "_chunk_relative_index": relative_idx,
-                }
-                items_yielded += 1
-                yield key, chunk.shard_url, image_data, metadata
-
-            absolute_idx += 1
-
-        logger.info(
-            f"WebDataset chunk {chunk.chunk_id}: yielded {items_yielded} items "
-            f"from ranges {unprocessed_ranges}"
-        )

caption_flow/utils/shard_tracker.py

@@ -1,83 +0,0 @@
-"""Shard tracking using CheckpointTracker base class."""
-
-from pathlib import Path
-from typing import Dict, Any, List, Set
-
-from .checkpoint_tracker import CheckpointTracker
-
-
-class ShardTracker(CheckpointTracker):
-    """Tracks shard processing progress."""
-
-    def __init__(self, checkpoint_path: Path):
-        """Initialize shard tracker with checkpoint file."""
-        self.completed_shards: Set[str] = set()
-        self.partial_shards: Dict[str, Dict[str, Any]] = {}
-        super().__init__(checkpoint_path)
-
-    def _get_default_state(self) -> Dict[str, Any]:
-        """Return default state structure for new checkpoints."""
-        return {"completed_shards": [], "partial_shards": {}}
-
-    def _deserialize_state(self, data: Dict[str, Any]) -> None:
-        """Deserialize loaded data into instance state."""
-        self.completed_shards = set(data.get("completed_shards", []))
-        self.partial_shards = data.get("partial_shards", {})
-
-    def _serialize_state(self) -> Dict[str, Any]:
-        """Serialize instance state for saving."""
-        return {
-            "completed_shards": list(self.completed_shards),
-            "partial_shards": self.partial_shards,
-        }
-
-    def mark_complete(self, shard_name: str) -> None:
-        """Mark a shard as complete."""
-        self.completed_shards.add(shard_name)
-        if shard_name in self.partial_shards:
-            del self.partial_shards[shard_name]
-        self.save()
-
-    def update_partial(self, shard_name: str, processed_keys: List[str]) -> None:
-        """Update partial progress for a shard."""
-        self.partial_shards[shard_name] = {"keys": processed_keys, "count": len(processed_keys)}
-        self.save()
-
-    def get_processed_keys(self, shard_name: str) -> Set[str]:
-        """Get set of processed keys for a shard."""
-        if shard_name in self.completed_shards:
-            return set()  # All done
-
-        if shard_name in self.partial_shards:
-            return set(self.partial_shards[shard_name].get("keys", []))
-
-        return set()
-
-    def is_complete(self, shard_name: str) -> bool:
-        """Check if a shard is complete."""
-        return shard_name in self.completed_shards
-
-    def get_remaining_shards(self, all_shards: List[str]) -> List[str]:
-        """Get list of shards that still need processing."""
-        remaining = []
-        for s in all_shards:
-            shard_name = Path(s).stem
-
-            if shard_name not in self.completed_shards:
-                remaining.append(s)
-
-        return remaining
-
-    def get_stats(self) -> Dict[str, Any]:
-        """Get shard tracking statistics."""
-        base_stats = super().get_stats()
-        base_stats.update(
-            {
-                "completed_shards": len(self.completed_shards),
-                "partial_shards": len(self.partial_shards),
-                "total_partial_keys": sum(
-                    len(data.get("keys", [])) for data in self.partial_shards.values()
-                ),
-            }
-        )
-        return base_stats