vesper-wizard 2.3.1 → 2.3.2

package/build/python/vesper/core/asset_downloader.py

@@ -1,679 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import json
-import mimetypes
-import os
-import shutil
-import tempfile
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional
-from urllib.parse import urlparse
-
-import aiohttp
-
-from vesper.core.download_recipe import get_download_recipe
-
-try:
-    import aiofiles
-except Exception:  # pragma: no cover
-    aiofiles = None
-
-try:
-    import webdataset as wds
-except Exception:  # pragma: no cover
-    wds = None
-
-
-IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff", ".tif", ".svg"}
-
-
-@dataclass
-class DownloadResult:
-    dataset_id: str
-    source: str
-    output_dir: str
-    downloaded_assets: int
-    failed_assets: int
-    errors_file: str
-    metadata_file: str
-    output_format: str
-
-
-class AssetDownloader:
-    def __init__(
-        self,
-        output_root: str,
-        workers: int = 8,
-        recipes_dir: Optional[str] = None,
-        progress_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None] | None]] = None,
-    ) -> None:
-        self.output_root = Path(output_root)
-        self.workers = max(1, min(workers, 32))
-        self.recipes_dir = recipes_dir
-        self.progress_callback = progress_callback
-
-    async def _emit(self, stage: str, payload: Dict[str, Any]) -> None:
-        if not self.progress_callback:
-            return
-        maybe = self.progress_callback(stage, payload)
-        if asyncio.iscoroutine(maybe):
-            await maybe
-
-    @staticmethod
-    def _hydrate_kaggle_credentials() -> None:
-        try:
-            from config import get_all  # type: ignore
-            keys = get_all() or {}
-        except Exception:
-            keys = {}
-
-        username = keys.get("kaggle_username") or os.getenv("KAGGLE_USERNAME")
-        key = keys.get("kaggle_key") or os.getenv("KAGGLE_KEY")
-
-        if username:
-            os.environ["KAGGLE_USERNAME"] = str(username)
-        if key:
-            os.environ["KAGGLE_KEY"] = str(key)
-
-        username = os.getenv("KAGGLE_USERNAME")
-        key = os.getenv("KAGGLE_KEY")
-        if not username or not key:
-            return
-
-        kaggle_dir = Path.home() / ".kaggle"
-        kaggle_file = kaggle_dir / "kaggle.json"
-        try:
-            kaggle_dir.mkdir(parents=True, exist_ok=True)
-            kaggle_file.write_text(
-                json.dumps({"username": username, "key": key}, ensure_ascii=False),
-                encoding="utf-8",
-            )
-            try:
-                os.chmod(kaggle_file, 0o600)
-            except Exception:
-                pass
-        except Exception:
-            pass
-
-    @staticmethod
-    def find_image_column(dataset: Any) -> Optional[str]:
-        """Auto-detect the image column in a HuggingFace dataset.
-        
-        Detection strategy (in priority order):
-        1. HF Feature type: columns with Image() feature type
-        2. Known column names: 'image', 'img', 'photo', 'image_url', etc.
-        3. URL pattern detection: columns containing image URLs (http(s)://...jpg)
-        4. Path pattern detection: columns with file paths ending in image extensions
-        """
-        # Strategy 1: Check HF Feature types (most reliable)
-        features = getattr(dataset, "features", None)
-        if features:
-            for name, feature in features.items():
-                feat_cls = feature.__class__.__name__.lower()
-                feat_str = str(feature).lower()
-                if feat_cls == "image" or "image(" in feat_str:
-                    return str(name)
-
-        # Strategy 2: Check known column names
-        cols = getattr(dataset, "column_names", []) or []
-        
-        # Exact match first (highest priority names)
-        priority_exact = ["image", "img", "photo", "picture", "images"]
-        for c in priority_exact:
-            if c in cols:
-                return c
-        
-        # Partial match (column names containing image-related keywords)
-        priority_partial = [
-            "image_path", "image_url", "img_path", "img_url",
-            "image_file", "file_name", "filepath", "filename",
-            "photo_url", "picture_url", "thumbnail",
-            "url", "path", "file",
-        ]
-        for target in priority_partial:
-            for c in cols:
-                if c.lower() == target:
-                    return c
-        
-        # Strategy 3: Sample values to detect URL/path patterns
-        try:
-            sample_size = min(5, len(dataset)) if hasattr(dataset, "__len__") else 5
-            if sample_size > 0:
-                for c in cols:
-                    is_image_col = False
-                    for i in range(sample_size):
-                        try:
-                            val = dataset[i][c]
-                        except Exception:
-                            break
-                        
-                        if val is None:
-                            continue
-                        
-                        # PIL Image object
-                        if hasattr(val, "save") and hasattr(val, "size"):
-                            is_image_col = True
-                            break
-                        
-                        # Dict with image data
-                        if isinstance(val, dict) and any(k in val for k in ("bytes", "path", "url")):
-                            is_image_col = True
-                            break
-                        
-                        # String: URL or file path
-                        if isinstance(val, str):
-                            val_lower = val.lower()
-                            # Check for image URLs
-                            if val_lower.startswith(("http://", "https://")) and any(
-                                ext in val_lower.split("?")[0] for ext in IMAGE_EXTENSIONS
-                            ):
-                                is_image_col = True
-                                break
-                            # Check for file paths with image extensions
-                            if any(val_lower.endswith(ext) for ext in IMAGE_EXTENSIONS):
-                                is_image_col = True
-                                break
-                    
-                    if is_image_col:
-                        return c
-        except Exception:
-            pass
-        
-        return None
-
-    async def download_assets(
-        self,
-        dataset_id: str,
-        source: Optional[str] = None,
-        repo_id: Optional[str] = None,
-        kaggle_ref: Optional[str] = None,
-        urls: Optional[List[str]] = None,
-        output_format: str = "webdataset",
-        output_dir: Optional[str] = None,
-        max_items: Optional[int] = None,
-        image_column: Optional[str] = None,
-    ) -> Dict[str, Any]:
-        recipe = get_download_recipe(dataset_id, self.recipes_dir)
-        if recipe:
-            source = source or recipe.get("source")
-            repo_id = repo_id or recipe.get("repo_id")
-            image_column = image_column or recipe.get("image_column")
-
-        source = (source or "").lower()
-        if source not in {"huggingface", "kaggle", "url"}:
-            raise ValueError("source must be one of: huggingface, kaggle, url")
-
-        # --- Validate imports and args BEFORE creating any directories ---
-        if source == "huggingface":
-            if not repo_id:
-                raise ValueError("repo_id is required for source=huggingface")
-            try:
-                from datasets import load_dataset as _ld  # noqa: F401
-            except Exception as e:
-                raise RuntimeError(
-                    f"datasets package is required for HuggingFace downloads. "
-                    f"Install with: pip install datasets. Details: {e}"
-                )
-        elif source == "kaggle":
-            ref = kaggle_ref or repo_id
-            if not ref:
-                raise ValueError("kaggle_ref is required for source=kaggle")
-            try:
-                from kaggle.api.kaggle_api_extended import KaggleApi as _Ka  # noqa: F401
-            except Exception as e:
-                raise RuntimeError(
-                    f"kaggle package is required for Kaggle downloads. "
-                    f"Install with: pip install kaggle. Details: {e}"
-                )
-        else:
-            if not urls:
-                raise ValueError("urls are required for source=url")
-
-        # --- Now safe to create directories ---
-        if output_dir:
-            dataset_dir = Path(output_dir).expanduser().resolve()
-        else:
-            dataset_dir = self.output_root / dataset_id.replace("/", "_").replace(":", "_")
-        images_dir = dataset_dir / "images"
-        dataset_dir.mkdir(parents=True, exist_ok=True)
-        images_dir.mkdir(parents=True, exist_ok=True)
-
-        errors_file = dataset_dir / "errors.jsonl"
-        metadata_file = dataset_dir / "metadata.jsonl"
-
-        try:
-            if source == "huggingface":
-                summary = await self._download_huggingface(repo_id, dataset_id, images_dir, metadata_file, errors_file, max_items, image_column)
-            elif source == "kaggle":
-                ref = kaggle_ref or repo_id
-                summary = await self._download_kaggle(ref, dataset_id, images_dir, metadata_file, errors_file, max_items)
-            else:
-                summary = await self._download_urls(urls, dataset_id, images_dir, metadata_file, errors_file, max_items)
-        except Exception:
-            # Clean up empty directories on failure so we don't leave ghost artifacts
-            if images_dir.exists() and not any(images_dir.iterdir()):
-                shutil.rmtree(dataset_dir, ignore_errors=True)
-            raise
-
-        if output_format == "webdataset":
-            await self._write_webdataset(dataset_dir, images_dir, metadata_file)
-        elif output_format == "parquet":
-            await self._write_parquet(dataset_dir, metadata_file)
-
-        result = DownloadResult(
-            dataset_id=dataset_id,
-            source=source,
-            output_dir=str(dataset_dir),
-            downloaded_assets=summary["downloaded"],
-            failed_assets=summary["failed"],
-            errors_file=str(errors_file),
-            metadata_file=str(metadata_file),
-            output_format=output_format,
-        )
-        return result.__dict__
-
-    async def _download_huggingface(
-        self,
-        repo_id: str,
-        dataset_id: str,
-        images_dir: Path,
-        metadata_file: Path,
-        errors_file: Path,
-        max_items: Optional[int],
-        image_column: Optional[str],
-    ) -> Dict[str, int]:
-        from datasets import load_dataset  # validated in download_assets()
-        import warnings
-        warnings.filterwarnings("ignore", message=".*trust_remote_code.*")
-
-        await self._emit("start", {"source": "huggingface", "repo_id": repo_id})
-
-        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN") or None
-
-        # Try loading with multiple strategies
-        ds = None
-        load_errors = []
-        
-        for trust_rc in [True, False]:
-            for split_name in ["train", "test", "validation"]:
-                try:
-                    kwargs = {"path": repo_id, "split": split_name}
-                    if trust_rc:
-                        kwargs["trust_remote_code"] = True 
-                    if token:
-                        kwargs["token"] = token
-                    ds = load_dataset(**kwargs)
-                    break
-                except Exception as e:
-                    msg = str(e)
-                    # Immediately raise auth errors
-                    if any(x in msg for x in ["401", "403", "gated", "Unauthorized"]):
-                        raise RuntimeError(
-                            f"Authentication required for '{repo_id}'. "
-                            "This dataset may be gated or private. "
-                            "Use the configure_keys tool to set HF_TOKEN, then retry."
-                        )
-                    load_errors.append(msg)
-                    continue
-            if ds is not None:
-                break
-        
-        # Fallback: load without split
-        if ds is None:
-            try:
-                kwargs = {"path": repo_id, "trust_remote_code": True}
-                if token:
-                    kwargs["token"] = token
-                dd = load_dataset(**kwargs)
-                from datasets import DatasetDict
-                if isinstance(dd, DatasetDict):
-                    first_split = list(dd.keys())[0]
-                    ds = dd[first_split]
-                else:
-                    ds = dd
-            except Exception as e:
-                msg = str(e)
-                if any(x in msg for x in ["401", "403", "gated", "Unauthorized"]):
-                    raise RuntimeError(
-                        f"Authentication required for '{repo_id}'. "
-                        "Use the configure_keys tool to set HF_TOKEN, then retry."
-                    )
-                combined = "; ".join(load_errors[:3])
-                raise RuntimeError(
-                    f"Failed to load HuggingFace dataset '{repo_id}': {msg}. "
-                    f"Previous attempts: {combined}"
-                )
-
-        col = image_column or self.find_image_column(ds)
-        if not col:
-            raise RuntimeError(
-                f"No image column detected in HuggingFace dataset '{repo_id}'. "
-                "Available columns: " + ", ".join(getattr(ds, "column_names", [])) + ". "
-                "Provide image_column parameter explicitly."
-            )
-
-        total = len(ds) if hasattr(ds, "__len__") else 0
-        target = min(total, max_items) if max_items and total else (max_items or total or 0)
-
-        downloaded = 0
-        failed = 0
-
-        # Create an aiohttp session for URL-based images
-        session = None
-        
-        try:
-            with metadata_file.open("w", encoding="utf-8") as mf, errors_file.open("w", encoding="utf-8") as ef:
-                for idx, row in enumerate(ds):
-                    if max_items and idx >= max_items:
-                        break
-                    try:
-                        out_name = f"{idx:08d}.jpg"
-                        out_path = images_dir / out_name
-                        value = row.get(col)
-                        
-                        # Handle URL-based images inline
-                        if isinstance(value, dict) and value.get("url") and not value.get("bytes") and not value.get("path"):
-                            url = value["url"]
-                            if session is None:
-                                session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60))
-                            await self._download_image_from_url(session, url, out_path)
-                        elif isinstance(value, str) and value.startswith(("http://", "https://")):
-                            if session is None:
-                                session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60))
-                            await self._download_image_from_url(session, value, out_path)
-                        else:
-                            self._save_image_value(value, out_path)
-
-                        record = {
-                            "dataset_id": dataset_id,
-                            "index": idx,
-                            "image_path": str(out_path),
-                            "source": "huggingface",
-                            "repo_id": repo_id,
-                        }
-                        mf.write(json.dumps(record, ensure_ascii=False) + "\n")
-                        downloaded += 1
-                        if downloaded % 50 == 0:
-                            await self._emit("progress", {"downloaded": downloaded, "failed": failed, "target": target})
-                    except Exception as e:
-                        failed += 1
-                        ef.write(json.dumps({"index": idx, "error": str(e)}, ensure_ascii=False) + "\n")
-        finally:
-            if session is not None:
-                await session.close()
-
-        await self._emit("done", {"downloaded": downloaded, "failed": failed})
-        return {"downloaded": downloaded, "failed": failed}
-
-    async def _download_image_from_url(self, session: aiohttp.ClientSession, url: str, out_path: Path) -> None:
-        """Download an image from a URL to a local path."""
-        async with session.get(url) as response:
-            if response.status != 200:
-                raise RuntimeError(f"HTTP {response.status} downloading {url}")
-            data = await response.read()
-            if not data:
-                raise RuntimeError(f"Empty response from {url}")
-            out_path.write_bytes(data)
-
-    async def _download_kaggle(
-        self,
-        kaggle_ref: str,
-        dataset_id: str,
-        images_dir: Path,
-        metadata_file: Path,
-        errors_file: Path,
-        max_items: Optional[int],
-    ) -> Dict[str, int]:
-        from kaggle.api.kaggle_api_extended import KaggleApi  # validated in download_assets()
-
-        await self._emit("start", {"source": "kaggle", "dataset": kaggle_ref})
-
-        self._hydrate_kaggle_credentials()
-
-        api = KaggleApi()
-        try:
-            api.authenticate()
-        except Exception as e:
-            raise RuntimeError(
-                "Kaggle authentication failed. Run 'configure_kaggle' or 'configure_keys' with "
-                "kaggle_username and kaggle_key, then retry. "
-                f"Details: {e}"
-            )
-
-        tmp_dir = Path(tempfile.mkdtemp(prefix="vesper_kaggle_assets_"))
-        downloaded = 0
-        failed = 0
-
-        try:
-            api.dataset_download_files(kaggle_ref, path=str(tmp_dir), unzip=True, quiet=True)
-            candidates = [p for p in tmp_dir.rglob("*") if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS]
-            if max_items:
-                candidates = candidates[:max_items]
-
-            with metadata_file.open("w", encoding="utf-8") as mf, errors_file.open("w", encoding="utf-8") as ef:
-                for idx, src_path in enumerate(candidates):
-                    try:
-                        out_name = f"{idx:08d}{src_path.suffix.lower()}"
-                        out_path = images_dir / out_name
-                        shutil.copy2(src_path, out_path)
-                        record = {
-                            "dataset_id": dataset_id,
-                            "index": idx,
-                            "image_path": str(out_path),
-                            "source": "kaggle",
-                            "repo_id": kaggle_ref,
-                        }
-                        mf.write(json.dumps(record, ensure_ascii=False) + "\n")
-                        downloaded += 1
-                    except Exception as e:
-                        failed += 1
-                        ef.write(json.dumps({"file": str(src_path), "error": str(e)}, ensure_ascii=False) + "\n")
-        finally:
-            shutil.rmtree(tmp_dir, ignore_errors=True)
-
-        await self._emit("done", {"downloaded": downloaded, "failed": failed})
-        return {"downloaded": downloaded, "failed": failed}
-
-    async def _download_urls(
-        self,
-        urls: List[str],
-        dataset_id: str,
-        images_dir: Path,
-        metadata_file: Path,
-        errors_file: Path,
-        max_items: Optional[int],
-    ) -> Dict[str, int]:
-        if aiofiles is None:
-            raise RuntimeError("aiofiles is required for URL downloads. Install with: pip install aiofiles")
-
-        selected = urls[:max_items] if max_items else urls
-        sem = asyncio.Semaphore(self.workers)
-
-        downloaded = 0
-        failed = 0
-        metadata_lock = asyncio.Lock()
-
-        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=180)) as session:
-            async def worker(idx: int, url: str) -> None:
-                nonlocal downloaded, failed
-                async with sem:
-                    try:
-                        local_path = await self._download_one_url(session, idx, url, images_dir)
-                        async with metadata_lock:
-                            async with aiofiles.open(metadata_file, "a", encoding="utf-8") as mf:
-                                await mf.write(json.dumps({
-                                    "dataset_id": dataset_id,
-                                    "index": idx,
-                                    "image_path": str(local_path),
-                                    "source": "url",
-                                    "url": url,
-                                }, ensure_ascii=False) + "\n")
-                        downloaded += 1
-                    except Exception as e:
-                        failed += 1
-                        async with metadata_lock:
-                            async with aiofiles.open(errors_file, "a", encoding="utf-8") as ef:
-                                await ef.write(json.dumps({"index": idx, "url": url, "error": str(e)}, ensure_ascii=False) + "\n")
-
-            tasks = [asyncio.create_task(worker(i, u)) for i, u in enumerate(selected)]
-            await asyncio.gather(*tasks)
-
-        await self._emit("done", {"downloaded": downloaded, "failed": failed})
-        return {"downloaded": downloaded, "failed": failed}
-
-    async def _download_one_url(self, session: aiohttp.ClientSession, idx: int, url: str, images_dir: Path) -> Path:
-        ext = Path(url.split("?")[0]).suffix.lower()
-        if ext not in IMAGE_EXTENSIONS:
-            ext = ".jpg"
-        out_path = images_dir / f"{idx:08d}{ext}"
-
-        existing_size = out_path.stat().st_size if out_path.exists() else 0
-        headers: Dict[str, str] = {}
-        if existing_size > 0:
-            headers["Range"] = f"bytes={existing_size}-"
-
-        async with session.get(url, headers=headers) as response:
-            if response.status not in (200, 206):
-                raise RuntimeError(f"HTTP {response.status}")
-
-            mode = "ab" if response.status == 206 and existing_size > 0 else "wb"
-            async with aiofiles.open(out_path, mode) as f:
-                async for chunk in response.content.iter_chunked(1024 * 256):
-                    await f.write(chunk)
-
-        return out_path
-
-    @staticmethod
-    def _save_image_value(value: Any, out_path: Path) -> None:
-        """Save an image value to disk. Handles multiple image representations:
-        - PIL Image objects (have .save method)
-        - dict with 'bytes' key (raw image bytes)
-        - dict with 'path' key (local file path)
-        - bytes/bytearray (raw image data)
-        - str (local file path)
-        """
-        if value is None:
-            raise ValueError("empty image value")
-
-        # PIL Image object
-        if hasattr(value, "save") and hasattr(value, "size"):
-            value.save(out_path)
-            return
-
-        # Raw bytes
-        if isinstance(value, (bytes, bytearray)):
-            out_path.write_bytes(value)
-            return
-
-        # Dict with image data
-        if isinstance(value, dict):
-            if value.get("bytes"):
-                raw = value["bytes"]
-                if isinstance(raw, (bytes, bytearray)):
-                    out_path.write_bytes(raw)
-                else:
-                    # Could be a list of ints
-                    out_path.write_bytes(bytes(raw))
-                return
-            if value.get("path"):
-                p = str(value["path"])
-                if os.path.exists(p):
-                    shutil.copy2(p, out_path)
-                    return
-                raise ValueError(f"Image path not found: {p}")
-            if value.get("url"):
-                raise ValueError("image URL detected — use async URL downloader")
-
-        # String: local file path
-        if isinstance(value, str):
-            if os.path.exists(value):
-                shutil.copy2(value, out_path)
-                return
-            if value.startswith(("http://", "https://")):
-                raise ValueError("image URL detected — use async URL downloader")
-            raise ValueError(f"Image path not found: {value}")
-
-        # numpy array (common in some datasets)
-        try:
-            import numpy as np
-            if isinstance(value, np.ndarray):
-                from PIL import Image
-                img = Image.fromarray(value)
-                img.save(out_path)
-                return
-        except (ImportError, Exception):
-            pass
-
-        raise ValueError(f"Unsupported image value type: {type(value).__name__}")
-
-    async def _write_webdataset(self, dataset_dir: Path, images_dir: Path, metadata_file: Path) -> None:
-        """Write a webdataset-compatible tar archive.
-
-        Uses Python's built-in tarfile module instead of wds.ShardWriter to
-        avoid the gopen() handler issue on Windows (backslash paths).
-        The resulting .tar files are fully compatible with webdataset readers.
-        """
-        import io
-        import tarfile as _tarfile
-
-        max_per_shard = 5000
-        shard_idx = 0
-        count_in_shard = 0
-        current_tar: _tarfile.TarFile | None = None
-
-        def _open_shard() -> _tarfile.TarFile:
-            nonlocal shard_idx
-            shard_path = dataset_dir / f"shard-{shard_idx:06d}.tar"
-            shard_idx += 1
-            return _tarfile.open(str(shard_path), "w")
-
-        try:
-            current_tar = _open_shard()
-
-            with metadata_file.open("r", encoding="utf-8") as mf:
-                for line in mf:
-                    row = json.loads(line)
-                    image_path = Path(row["image_path"])
-                    if not image_path.exists():
-                        continue
-
-                    key = image_path.stem
-                    ext = image_path.suffix.lstrip(".") or "jpg"
-
-                    # Add image file
-                    img_data = image_path.read_bytes()
-                    img_info = _tarfile.TarInfo(name=f"{key}.{ext}")
-                    img_info.size = len(img_data)
-                    current_tar.addfile(img_info, io.BytesIO(img_data))
-
-                    # Add JSON metadata sidecar
-                    json_data = json.dumps(row, ensure_ascii=False).encode("utf-8")
-                    json_info = _tarfile.TarInfo(name=f"{key}.json")
-                    json_info.size = len(json_data)
-                    current_tar.addfile(json_info, io.BytesIO(json_data))
-
-                    count_in_shard += 1
-                    if count_in_shard >= max_per_shard:
-                        current_tar.close()
-                        current_tar = _open_shard()
-                        count_in_shard = 0
-        finally:
-            if current_tar is not None:
-                current_tar.close()
-
-    async def _write_parquet(self, dataset_dir: Path, metadata_file: Path) -> None:
-        try:
-            import pyarrow as pa
-            import pyarrow.parquet as pq
-        except Exception as e:
-            raise RuntimeError(f"pyarrow is required for parquet output: {e}")
-
-        rows: List[Dict[str, Any]] = []
-        with metadata_file.open("r", encoding="utf-8") as mf:
-            for line in mf:
-                rows.append(json.loads(line))
-
-        table = pa.Table.from_pylist(rows)
-        pq.write_table(table, str(dataset_dir / "metadata.parquet"))