vesper-wizard 2.0.5 → 2.0.6

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

@@ -0,0 +1,675 @@
+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",
+        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 ---
+        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"))