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

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

caption_flow-0.2.3.dist-info/RECORD

@@ -1,35 +0,0 @@
-caption_flow/__init__.py,sha256=NLPJ25lRN7xHqncXweINDNwbt0q8lgjZ30G21zlPdRs,303
-caption_flow/cli.py,sha256=qEueeJhf3DvxSBxnOp5t32p6gAnZskvIDe6cwtPA0-Y,28892
-caption_flow/models.py,sha256=bpr7yMy3vPErZCQwmgOYIix489rRGbT6lVw8wxxwTkc,4931
-caption_flow/monitor.py,sha256=bAt9EJqfPgT_KdbknGdCxwBRH002pRDgyUmYIj6Dyso,7885
-caption_flow/orchestrator.py,sha256=ciqWghxUxk-5s6u7W3JwD7_JLSFYV57NgOwiMkxME-I,36133
-caption_flow/storage.py,sha256=Wqgtsk6yZ9Kf-izeUKHLwSvPUH3xFqIbzox20QHbc64,43370
-caption_flow/processors/__init__.py,sha256=hvq-OuAJWQe6hFglKe7QmkS8473k20FmxZDSxfXpCrg,423
-caption_flow/processors/base.py,sha256=JlTqCHo5HRXrXMVzgle_6pNwh4HGHsF7jLF6PeSnWr0,6783
-caption_flow/processors/huggingface.py,sha256=MNz9vDMtrrTOSXe9Q_kbBrQ7XBv69X6x5xD_QP9icdg,33765
-caption_flow/processors/local_filesystem.py,sha256=EYmsImbkqsIU7UZL2FijL0hotKLtPOtkzfwernQDSxA,27860
-caption_flow/processors/webdataset.py,sha256=xsrYx7_5FCqez30dc4hSDYfyA9A0oKqHqwt7CRc1J0c,33812
-caption_flow/utils/__init__.py,sha256=F1BChVoCsj9zn1GJRBOLHET1kLW6xrAmsbzcR7hHy6Y,202
-caption_flow/utils/auth.py,sha256=UrxX2n8OEEcfMD1Ey27TxGfrJFmUCpC59x-SCrQJoVE,2253
-caption_flow/utils/caption_utils.py,sha256=esUMAdcCkNjRroZ0Bhxv0_yKlLtMf0XeDCTt-5k6bik,5309
-caption_flow/utils/certificates.py,sha256=eu4blQZEkL9NRaY1ynQWg1asvDorRYhGRZea7STonJE,4635
-caption_flow/utils/checkpoint_tracker.py,sha256=-nN5gLvXyMdKOCT2SNNL2Km6UYm2Hii9wuXeezWhwx4,3339
-caption_flow/utils/chunk_tracker.py,sha256=x9UwFxpj-nMeAJ6bpKw5E09QNUqu7L0pejTlk8nxgE8,19402
-caption_flow/utils/dataset_loader.py,sha256=2-SgXPGQkF4CyA3zyVYfSbZMSk4YzTsVFY0izmOZPrM,8771
-caption_flow/utils/dataset_metadata_cache.py,sha256=AJ8Z1GYT0DC9_LLjxNvrePKU7ecenNZun5GhaB2gvj0,2650
-caption_flow/utils/image_processor.py,sha256=7Ed92iUJ-OvjzQmAGPaULoYEqoirVHHo0lxtceWGc44,5586
-caption_flow/utils/job_queue.py,sha256=itdfXcrkvGjmXn4qtpgMF63k1ufRBaejDe4V6WcxzgU,1104
-caption_flow/utils/json_utils.py,sha256=IiZYn8uCM-3pYmyIbX2fmaOIyutArn67SqAyp0ggNpU,5396
-caption_flow/utils/prompt_template.py,sha256=AKp0diSZqNBMwZkpiTNjw8-bbQwHStr7QZTOJ7o1dC4,4345
-caption_flow/utils/shard_processor.py,sha256=_PCW5TfSHFfCc63Sn7bVzgjA625-aWzL4cWwZLjW0rQ,3935
-caption_flow/utils/shard_tracker.py,sha256=1OqiueaC8WoxhY2nc03erZAc50mnQCZazATS6R14lbQ,3029
-caption_flow/utils/vllm_config.py,sha256=TC7Rmjk0zRKbBXbWUXrFL4Z58hzax_-4L0pXZn09hdM,6019
-caption_flow/workers/base.py,sha256=2AGWERC5hbmO-0V_A1MUbgRVvRNN3blqGPyDokvvzmM,7575
-caption_flow/workers/caption.py,sha256=_uvpdoBzym1TKWKXtky7hBfj8YnG1EaJz-NRwaH2X1A,36722
-caption_flow/workers/data.py,sha256=0Tg8NE0wdONeMlivYQ4nvbcfWdLuU51O7vR8_YSnJgo,14813
-caption_flow-0.2.3.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-caption_flow-0.2.3.dist-info/METADATA,sha256=bk5Gk3eWuDH_UWXPEDKulksPc3hVHvnzm3sstLbuU-0,11914
-caption_flow-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-caption_flow-0.2.3.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
-caption_flow-0.2.3.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
-caption_flow-0.2.3.dist-info/RECORD,,