@vespermcp/mcp-server 1.2.20 → 1.2.21

package/src/python/hf_fallback.py

@@ -5,6 +5,12 @@ Used when the HF Hub file listing finds no suitable data files
 (e.g. script-based datasets, gated datasets, datasets that use
 the `datasets` library format).
 
+Handles:
+- Legacy script-based datasets (trust_remote_code)
+- Gated/private datasets (token auth)
+- Image datasets (PIL Image columns → stripped for tabular export)
+- Various split formats (DatasetDict, single split)
+
 Usage:
     python hf_fallback.py '{"repo_id": "user/dataset", "output_path": "/path/to/output.parquet", "token": "optional", "max_rows": 100000, "split": "train"}'
 
@@ -15,6 +21,127 @@ Output: JSON to stdout
 import sys
 import json
 import os
+import warnings
+
+# Suppress noisy HF warnings about trust_remote_code etc.
+warnings.filterwarnings("ignore", message=".*trust_remote_code.*")
+warnings.filterwarnings("ignore", message=".*legacy.*")
+
+
+def _detect_image_columns(ds):
+    """Detect columns that contain HF Image features or PIL Image objects."""
+    image_cols = []
+    features = getattr(ds, "features", None)
+    if features:
+        for name, feat in features.items():
+            feat_cls = feat.__class__.__name__.lower()
+            feat_str = str(feat).lower()
+            if feat_cls == "image" or "image(" in feat_str:
+                image_cols.append(name)
+    return image_cols
+
+
+def _strip_image_columns(ds, image_cols):
+    """Remove image columns from dataset so it can be exported to Parquet/CSV.
+    
+    Image columns contain PIL Image objects that can't be serialized to tabular
+    formats. We replace them with a placeholder string indicating the column
+    was an image column.
+    """
+    if not image_cols:
+        return ds
+    
+    # Remove the image columns entirely for tabular export
+    cols_to_keep = [c for c in ds.column_names if c not in image_cols]
+    if not cols_to_keep:
+        # Dataset is ALL image columns — keep them but cast to path strings if possible
+        return ds
+    
+    return ds.select_columns(cols_to_keep)
+
+
+def _load_dataset_robust(repo_id, token, split):
+    """Load a HuggingFace dataset with multiple fallback strategies.
+    
+    Strategy order:
+    1. Normal load with trust_remote_code=True (handles legacy script datasets)
+    2. Load without trust_remote_code (newer datasets that reject it)
+    3. Load with streaming=True then materialize (handles very large datasets)
+    """
+    from datasets import load_dataset, DatasetDict
+
+    errors = []
+    splits_to_try = [split] if split else ["train", "test", "validation", None]
+    
+    # Strategy 1: Normal load with trust_remote_code
+    for s in splits_to_try:
+        try:
+            kwargs = {"path": repo_id, "trust_remote_code": True}
+            if token:
+                kwargs["token"] = token
+            if s:
+                kwargs["split"] = s
+            ds = load_dataset(**kwargs)
+            return ds, s
+        except (ValueError, KeyError):
+            continue
+        except Exception as e:
+            msg = str(e)
+            # Auth errors should be raised immediately, not retried
+            if any(x in msg for x in ["401", "403", "gated", "Unauthorized", "access"]):
+                raise
+            if "split" in msg.lower() or "key" in msg.lower():
+                continue
+            errors.append(f"trust_remote_code=True, split={s}: {msg}")
+    
+    # Strategy 2: Load WITHOUT trust_remote_code (some repos reject it)
+    for s in splits_to_try:
+        try:
+            kwargs = {"path": repo_id}
+            if token:
+                kwargs["token"] = token
+            if s:
+                kwargs["split"] = s
+            ds = load_dataset(**kwargs)
+            return ds, s
+        except (ValueError, KeyError):
+            continue
+        except Exception as e:
+            msg = str(e)
+            if any(x in msg for x in ["401", "403", "gated", "Unauthorized", "access"]):
+                raise
+            if "split" in msg.lower() or "key" in msg.lower():
+                continue
+            errors.append(f"trust_remote_code=False, split={s}: {msg}")
+    
+    # Strategy 3: Streaming fallback (for very large / oddly structured datasets)
+    for s in splits_to_try:
+        if s is None:
+            continue  # streaming requires a split
+        try:
+            kwargs = {"path": repo_id, "streaming": True, "trust_remote_code": True}
+            if token:
+                kwargs["token"] = token
+            if s:
+                kwargs["split"] = s
+            ds_stream = load_dataset(**kwargs)
+            # Materialize from streaming iterator
+            from datasets import Dataset as HFDataset
+            rows = []
+            for i, row in enumerate(ds_stream):
+                if i >= 500000:
+                    break
+                rows.append(row)
+            if rows:
+                ds = HFDataset.from_list(rows)
+                return ds, s
+        except Exception:
+            continue
+    
+    # All strategies failed
+    error_summary = "; ".join(errors[:3]) if errors else "No valid configuration found"
+    return None, error_summary
+
 
 def main():
     if len(sys.argv) < 2:
@@ -53,37 +180,10 @@ def main():
         pl = None
 
     try:
-        # Try loading with streaming first (memory-efficient)
-        # If split is not specified, try common ones
-        splits_to_try = [split] if split else ["train", "test", "validation", None]
-        
-        ds = None
-        used_split = None
-        
-        for s in splits_to_try:
-            try:
-                kwargs = {
-                    "path": repo_id,
-                    "trust_remote_code": True,
-                }
-                if token:
-                    kwargs["token"] = token
-                if s:
-                    kwargs["split"] = s
-
-                ds = load_dataset(**kwargs)
-                used_split = s
-                break
-            except (ValueError, KeyError):
-                # Split doesn't exist, try next
-                continue
-            except Exception as e:
-                if "split" in str(e).lower() or "key" in str(e).lower():
-                    continue
-                raise
+        ds, used_split = _load_dataset_robust(repo_id, token, split)
 
         if ds is None:
-            print(json.dumps({"ok": False, "error": f"Could not load dataset '{repo_id}'. No valid splits found."}))
+            print(json.dumps({"ok": False, "error": f"Could not load dataset '{repo_id}': {used_split}"}))
             sys.exit(1)
 
         # Handle DatasetDict (when no split specified)
@@ -107,37 +207,88 @@ def main():
             ds = ds.select(range(max_rows))
             total_rows = max_rows
 
+        # Detect and handle image columns (PIL Image objects can't be exported to Parquet)
+        image_cols = _detect_image_columns(ds)
+        has_images = len(image_cols) > 0
+        
+        if has_images:
+            # Strip image columns for tabular export, note them in output
+            export_ds = _strip_image_columns(ds, image_cols)
+        else:
+            export_ds = ds
+
         # Ensure output directory exists
         os.makedirs(os.path.dirname(output_path), exist_ok=True)
 
         # Export to parquet
-        columns = ds.column_names
+        columns = export_ds.column_names
 
-        if output_path.endswith(".parquet"):
-            ds.to_parquet(output_path)
-        elif output_path.endswith(".csv"):
-            ds.to_csv(output_path)
-        else:
-            # Default to parquet
-            if not output_path.endswith(".parquet"):
-                output_path = output_path + ".parquet"
-            ds.to_parquet(output_path)
+        try:
+            if output_path.endswith(".parquet"):
+                export_ds.to_parquet(output_path)
+            elif output_path.endswith(".csv"):
+                export_ds.to_csv(output_path)
+            else:
+                # Default to parquet
+                if not output_path.endswith(".parquet"):
+                    output_path = output_path + ".parquet"
+                export_ds.to_parquet(output_path)
+        except Exception as export_err:
+            # If parquet export fails (e.g. complex nested types), try CSV
+            csv_path = output_path.replace(".parquet", ".csv")
+            try:
+                export_ds.to_csv(csv_path)
+                output_path = csv_path
+            except Exception:
+                raise export_err  # Re-raise original error
 
-        print(json.dumps({
+        result = {
             "ok": True,
             "path": output_path,
             "rows": total_rows,
             "columns": columns,
             "split": used_split
-        }))
+        }
+        
+        if has_images:
+            result["image_columns"] = image_cols
+            result["note"] = (
+                f"This dataset contains image columns ({', '.join(image_cols)}). "
+                "Image data was stripped for tabular export. "
+                "Use vesper_download_assets with source='huggingface' to download the actual images."
+            )
+
+        print(json.dumps(result))
 
     except Exception as e:
         error_msg = str(e)
-        # Provide helpful hints
-        if "401" in error_msg or "403" in error_msg or "gated" in error_msg.lower():
-            error_msg += " (This dataset may be gated/private. Set HF_TOKEN via configure_keys tool.)"
-        elif "FileNotFoundError" in error_msg or "does not exist" in error_msg.lower():
+        # Provide helpful, actionable hints
+        if "401" in error_msg or "Unauthorized" in error_msg:
+            error_msg = (
+                f"Authentication required for dataset '{repo_id}'. "
+                "This dataset may be gated or private. "
+                "Use the configure_keys tool to set your HF_TOKEN, then retry."
+            )
+        elif "403" in error_msg or "Forbidden" in error_msg:
+            error_msg = (
+                f"Access denied for dataset '{repo_id}'. "
+                "You may need to accept the dataset's usage agreement on huggingface.co, "
+                "then set HF_TOKEN via configure_keys tool."
+            )
+        elif "gated" in error_msg.lower():
+            error_msg = (
+                f"Dataset '{repo_id}' is gated. "
+                "Visit https://huggingface.co/datasets/{repo_id} to request access, "
+                "then set HF_TOKEN via configure_keys tool."
+            ).format(repo_id=repo_id)
+        elif "FileNotFoundError" in error_msg or "does not exist" in error_msg.lower() or "doesn't exist" in error_msg.lower():
             error_msg = f"Dataset '{repo_id}' not found on HuggingFace. Check the dataset ID."
+        elif "script" in error_msg.lower() and "no longer supported" in error_msg.lower():
+            error_msg = (
+                f"Dataset '{repo_id}' uses a legacy loading script that is no longer supported "
+                "by the current version of the datasets library. "
+                "Try: pip install datasets --upgrade, or use an older datasets version."
+            )
 
         print(json.dumps({"ok": False, "error": error_msg}))
         sys.exit(1)

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

@@ -9,6 +9,7 @@ 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
 
@@ -25,7 +26,7 @@ except Exception:  # pragma: no cover
     wds = None
 
 
-IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff"}
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff", ".tif", ".svg"}
 
 
 @dataclass
@@ -98,22 +99,88 @@ class AssetDownloader:
 
     @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():
-                feature_name = feature.__class__.__name__.lower()
-                feature_repr = str(feature).lower()
-                if feature_name == "image" or "image(" in feature_repr:
-                    return str(name)
-                lower = str(name).lower()
-                if lower in {"image", "images", "img", "image_path", "image_url", "url"}:
+                feat_cls = feature.__class__.__name__.lower()
+                feat_str = str(feature).lower()
+                if feat_cls == "image" or "image(" in feat_str:
                     return str(name)
 
-        candidate_columns = ["image", "images", "img", "image_path", "image_url", "url", "file_name", "filepath"]
+        # Strategy 2: Check known column names
         cols = getattr(dataset, "column_names", []) or []
-        for c in candidate_columns:
+        
+        # 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(
@@ -214,20 +281,73 @@ class AssetDownloader:
         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})
 
-        try:
-            ds = load_dataset(repo_id, split="train")
-        except Exception:
-            dd = load_dataset(repo_id)
-            first_split = list(dd.keys())[0]
-            ds = dd[first_split]
+        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(
-                "No image column detected in HuggingFace dataset. Provide image_column or use fallback strategy with URL column."
+                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
@@ -236,33 +356,63 @@ class AssetDownloader:
         downloaded = 0
         failed = 0
 
-        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
-                    self._save_image_value(row.get(col), 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")
+        # 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,
@@ -393,30 +543,66 @@ class AssetDownloader:
 
     @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")
 
-        if hasattr(value, "save"):
+        # 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"):
-                out_path.write_bytes(value["bytes"])
-                return
-            if value.get("path") and os.path.exists(value["path"]):
-                shutil.copy2(value["path"], out_path)
+                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 requires URL downloader fallback")
+                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
-            raise ValueError("string image value is not a local path")
+            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)}")
+        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.