diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index d21974d540..25965a847e 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -83,16 +83,12 @@ def _serialize_preview_rows(rows): # --- Endpoints --- # Recognized data-file extensions for the single-file fallback approach. -DATA_EXTS = ( - '.parquet', - '.json', '.jsonl', - '.csv', '.tsv', - '.txt', - '.arrow', - '.tar', '.tar.gz', '.tgz', - '.gz', '.zst', - '.zip', -) +# Tabular formats are preferred over archives for Tier 1 preview because +# archives (e.g. images.zip) may be loaded as ImageFolder datasets with +# synthetic columns (image/label) that don't match the real dataset schema. +_TABULAR_EXTS = ('.parquet', '.json', '.jsonl', '.csv', '.tsv', '.arrow') +_ARCHIVE_EXTS = ('.tar', '.tar.gz', '.tgz', '.gz', '.zst', '.zip', '.txt') +DATA_EXTS = _TABULAR_EXTS + _ARCHIVE_EXTS LOCAL_FILE_EXTS = ('.json', '.jsonl', '.csv', '.parquet') LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"} BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -363,8 +359,20 @@ def check_format( ) data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)] - if data_files: - first_file = data_files[0] + # Prefer tabular formats over archives (e.g. images.zip → ImageFolder + # with synthetic image/label columns that don't match the real schema). + tabular_files = [f for f in data_files if any(f.endswith(ext) for ext in _TABULAR_EXTS)] + candidates = tabular_files or data_files + + # When a subset is specified, narrow to files whose name matches + # (e.g. subset="testmini" → prefer "testmini.parquet"). + if request.subset and candidates: + subset_matches = [f for f in candidates if request.subset in Path(f).stem] + if subset_matches: + candidates = subset_matches + + if candidates: + first_file = candidates[0] logger.info(f"Tier 1: loading single file {first_file}") load_kwargs = { "path": request.dataset_name, diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index 4146006c18..5b1c832923 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/studio/backend/utils/datasets/__init__.py @@ -31,6 +31,7 @@ from .format_conversion import ( convert_alpaca_to_chatml, convert_to_vlm_format, convert_llava_to_vlm_format, + convert_sharegpt_with_images_to_vlm_format, ) # Chat templates @@ -81,6 +82,7 @@ __all__ = [ "convert_alpaca_to_chatml", "convert_to_vlm_format", "convert_llava_to_vlm_format", + "convert_sharegpt_with_images_to_vlm_format", # Templates "apply_chat_template_to_dataset", "get_dataset_info_summary", diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index b2a290355a..2b430cee14 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -31,6 +31,7 @@ from .format_conversion import ( convert_alpaca_to_chatml, convert_to_vlm_format, convert_llava_to_vlm_format, + convert_sharegpt_with_images_to_vlm_format, ) from .chat_templates import ( apply_chat_template_to_dataset, @@ -688,17 +689,15 @@ def format_and_template_dataset( "errors": [], } except Exception as e: - errors.append(f"Failed to apply user VLM mapping: {e}") - return { - "dataset": dataset, - "detected_format": "user_mapped", - "final_format": "vlm_conversion_failed", - "is_vlm": True, - "success": False, - "requires_manual_mapping": True, - "warnings": warnings, - "errors": errors, - } + # User mapping failed — fall back to auto-detection instead + # of giving up (handles stale cached mappings gracefully) + warnings.append( + f"User VLM mapping (image='{user_vlm_image_column}', " + f"text='{user_vlm_text_column}') failed: {e} — " + f"falling back to auto-detection" + ) + print(f"⚠️ User VLM mapping failed, falling back to auto-detection...") + custom_format_mapping = None # clear so auto-detection runs below else: errors.append( f"Invalid VLM mapping: need 'image' and 'text' roles. Got: {custom_format_mapping}" @@ -738,6 +737,33 @@ def format_and_template_dataset( "errors": errors, } + # Handle ShareGPT/ChatML + image column (e.g. ShareGPT4V, LLaVA-style) + elif vlm_structure["format"] == "sharegpt_with_images": + try: + dataset = convert_sharegpt_with_images_to_vlm_format( + dataset, + image_column=vlm_structure["image_column"], + messages_column=vlm_structure["messages_column"], + dataset_name=dataset_name, + progress_callback=progress_callback, + ) + warnings.append("Converted from ShareGPT+image format to standard VLM format") + except Exception as e: + errors.append(f"Failed to convert ShareGPT+image format: {e}") + import traceback + traceback.print_exc() + + return { + "dataset": dataset, + "detected_format": "sharegpt_with_images", + "final_format": "vlm_conversion_failed", + "is_vlm": True, + "success": False, + "requires_manual_mapping": True, + "warnings": warnings, + "errors": errors, + } + # Handle simple format elif vlm_structure["needs_conversion"]: if vlm_text_column is None: diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index bb80dd8d6e..7c45a03b49 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -8,6 +8,8 @@ This module contains functions for converting between dataset formats (Alpaca, ShareGPT, ChatML) and standardizing chat formats. """ +import os + from datasets import IterableDataset @@ -311,7 +313,7 @@ def convert_to_vlm_format( def _convert_single_sample(sample): """Convert a single sample to VLM format.""" - # Get image (might be PIL Image, local path, or URL) + # Get image (might be PIL Image, local path, URL, or bare filename) image_data = sample[image_column] if isinstance(image_data, str): @@ -320,11 +322,21 @@ def convert_to_vlm_format( from io import BytesIO with fsspec.open(image_data, "rb", expand=True) as f: image_data = Image.open(BytesIO(f.read())).convert("RGB") + elif _image_lookup is not None and image_data in _image_lookup: + # Bare filename → resolve via HF repo lookup + from huggingface_hub import hf_hub_download + local_path = hf_hub_download( + dataset_name, _image_lookup[image_data], repo_type="dataset", + ) + image_data = Image.open(local_path).convert("RGB") else: image_data = Image.open(image_data).convert("RGB") - # Get text + # Get text (if list of strings, pick a random one — e.g. multiple captions) text_data = sample[text_column] + if isinstance(text_data, list) and len(text_data) > 0: + import random + text_data = random.choice(text_data) # Get instruction (static or dynamic) if uses_dynamic and instruction_column: @@ -356,6 +368,36 @@ def convert_to_vlm_format( first_image = next(iter(dataset))[image_column] has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://")) + # ── Bare-filename detection: images stored as filenames (e.g. "img_001.png") + # that don't exist locally. Build a basename→repo_path lookup so we can + # resolve them via hf_hub_download during conversion. + _image_lookup = None + _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff') + if ( + not has_urls + and isinstance(first_image, str) + and not os.path.exists(first_image) + and dataset_name + ): + try: + from huggingface_hub import HfApi + _notify("Resolving image filenames from HF repo...") + print(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...") + repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset") + _image_lookup = { + os.path.basename(f): f + for f in repo_files + if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS) + } + if first_image in _image_lookup: + print(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')") + else: + print(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open") + _image_lookup = None + except Exception as e: + print(f"⚠️ Failed to build HF repo image lookup: {e}") + _image_lookup = None + # ── URL probe: 200 samples with parallel workers to estimate speed + failure rate ── PROBE_SIZE = 200 MAX_FAIL_RATE = 0.3 @@ -443,8 +485,10 @@ def convert_to_vlm_format( idx = futures[future] try: batch_results[idx] = future.result() - except Exception: + except Exception as e: failed_count += 1 + if failed_count == 1: + print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}") converted_list.extend(r for r in batch_results if r is not None) @@ -463,8 +507,11 @@ def convert_to_vlm_format( for sample in pbar: try: converted_list.append(_convert_single_sample(sample)) - except Exception: + except Exception as e: failed_count += 1 + if failed_count == 1: + # Log the first failure to aid debugging + print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}") pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False) pbar.close() @@ -494,6 +541,164 @@ def convert_to_vlm_format( return converted_list +def convert_sharegpt_with_images_to_vlm_format( + dataset, + image_column="image", + messages_column="conversations", + dataset_name=None, + progress_callback=None, +): + """ + Converts ShareGPT/ChatML datasets that have a separate image column and + ```` placeholders inside the conversation text. + + Example input:: + + { + "image": "sam/images/sa_545504.jpg", + "conversations": [ + {"from": "human", "value": "\\nWhat is this photo about?"}, + {"from": "gpt", "value": "The image captures..."} + ] + } + + Returns a list of dicts in standard VLM messages format (PIL Images inline). + """ + from PIL import Image + from tqdm import tqdm + + _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff') + _ROLE_MAP = { + "human": "user", "user": "user", + "gpt": "assistant", "assistant": "assistant", + "system": "system", + } + + def _notify(msg): + if progress_callback: + progress_callback(status_message=msg) + + # ── Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ── + total = len(dataset) + first_image = next(iter(dataset))[image_column] + + _image_lookup = None + if ( + isinstance(first_image, str) + and not first_image.startswith(("http://", "https://")) + and not os.path.exists(first_image) + and dataset_name + ): + try: + from huggingface_hub import HfApi + _notify("Resolving image filenames from HF repo...") + print(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...") + repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset") + _image_lookup = { + os.path.basename(f): f + for f in repo_files + if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS) + } + # Also add the full relative paths as keys (for paths like "sam/images/sa_545504.jpg") + for f in repo_files: + if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS): + _image_lookup[f] = f + if first_image in _image_lookup: + print(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')") + else: + print(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open") + _image_lookup = None + except Exception as e: + print(f"⚠️ Failed to build HF repo image lookup: {e}") + _image_lookup = None + + def _resolve_image(image_data): + """Resolve image data to a PIL Image object.""" + if hasattr(image_data, 'size') and hasattr(image_data, 'mode'): + return image_data # Already PIL + if isinstance(image_data, str): + if image_data.startswith(("http://", "https://")): + import fsspec + from io import BytesIO + with fsspec.open(image_data, "rb", expand=True) as f: + return Image.open(BytesIO(f.read())).convert("RGB") + elif _image_lookup is not None and image_data in _image_lookup: + from huggingface_hub import hf_hub_download + local_path = hf_hub_download( + dataset_name, _image_lookup[image_data], repo_type="dataset", + ) + return Image.open(local_path).convert("RGB") + else: + return Image.open(image_data).convert("RGB") + if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data): + if image_data.get("bytes"): + from io import BytesIO + return Image.open(BytesIO(image_data["bytes"])).convert("RGB") + if image_data.get("path"): + return Image.open(image_data["path"]).convert("RGB") + raise ValueError(f"Cannot resolve image: {type(image_data)}") + + def _convert_single_sample(sample): + """Convert a single ShareGPT+image sample to standard VLM format.""" + pil_image = _resolve_image(sample[image_column]) + conversation = sample[messages_column] + + new_messages = [] + for msg in conversation: + role_raw = msg.get("from") or msg.get("role", "user") + role = _ROLE_MAP.get(role_raw.lower(), role_raw.lower()) + text = msg.get("value") or msg.get("content") or "" + + # Split on to interleave text and image content blocks + if "" in text: + parts = text.split("") + content = [] + for i, part in enumerate(parts): + part = part.strip() + if part: + content.append({"type": "text", "text": part}) + if i < len(parts) - 1: + content.append({"type": "image", "image": pil_image}) + # If was the entire text, content might just be the image + if not content: + content.append({"type": "image", "image": pil_image}) + else: + content = [{"type": "text", "text": text}] + + new_messages.append({"role": role, "content": content}) + + return {"messages": new_messages} + + # ── Full conversion with progress ── + print(f"🔄 Converting {total} samples from ShareGPT+image format...") + converted_list = [] + failed_count = 0 + + pbar = tqdm(dataset, total=total, desc="Converting ShareGPT+image", unit="sample") + for sample in pbar: + try: + converted_list.append(_convert_single_sample(sample)) + except Exception as e: + failed_count += 1 + if failed_count == 1: + print(f"⚠️ First conversion failure: {type(e).__name__}: {e}") + pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False) + pbar.close() + + if failed_count > 0: + print(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples") + + if len(converted_list) == 0: + raise ValueError( + f"All {total} samples failed during ShareGPT+image conversion — " + "no usable samples found." + ) + + print(f"✅ Converted {len(converted_list)}/{total} samples") + _notify(f"Converted {len(converted_list):,}/{total:,} samples successfully") + return converted_list + + def convert_llava_to_vlm_format(dataset): """ Converts Llava format to standard VLM format. diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index b8385faec8..d44ca4c960 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -8,6 +8,13 @@ This module contains functions for detecting dataset formats (Alpaca, ShareGPT, detecting multimodal/VLM dataset structures, and heuristic-based column mapping. """ +import re + + +def _keyword_in_column(keyword: str, col_name: str) -> bool: + """Word-boundary keyword match to avoid false positives like 'pic' in 'topic'.""" + return re.search(r'\b' + re.escape(keyword) + r'\b', col_name, re.IGNORECASE) is not None + def detect_dataset_format(dataset): """ @@ -354,6 +361,7 @@ def detect_multimodal_dataset(dataset): 'image', 'img', 'pixel', 'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg', 'photo', 'pic', 'picture', 'visual', + 'file_name', 'filename', ] # Keywords that indicate audio data @@ -364,11 +372,11 @@ def detect_multimodal_dataset(dataset): modality_types = set() # ── Image detection ───────────────────────────────────── - # Pass 1: column-name heuristic + # Pass 1: column-name heuristic (word-boundary match to avoid + # false positives like 'pic' in 'topic') for col_name in column_names: - col_lower = col_name.lower() for keyword in image_keywords: - if keyword in col_lower: + if _keyword_in_column(keyword, col_name): multimodal_columns.append(col_name) modality_types.add(keyword) break @@ -384,11 +392,10 @@ def detect_multimodal_dataset(dataset): modality_types.add("image") # ── Audio detection ───────────────────────────────────── - # Pass 1: column-name heuristic + # Pass 1: column-name heuristic (word-boundary match) for col_name in column_names: - col_lower = col_name.lower() for keyword in audio_keywords: - if keyword in col_lower: + if _keyword_in_column(keyword, col_name): audio_columns.append(col_name) modality_types.add("audio") break @@ -471,6 +478,17 @@ def _is_image_value(value) -> bool: if isinstance(value, (bytes, bytearray)): return _has_image_header(value) + # String that looks like an image file path or URL + _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff', '.svg') + if isinstance(value, str) and len(value) < 1000: + lower = value.strip().lower() + # Image URL (http://... ending in image extension) + if lower.startswith(("http://", "https://")) and any(lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS): + return True + # Image file path (relative or absolute path ending in image extension) + if any(lower.endswith(ext) for ext in _IMAGE_EXTS): + return True + return False @@ -575,6 +593,46 @@ def detect_vlm_dataset_structure(dataset): "text_column": None, } + # Check for ShareGPT/ChatML conversations with placeholder + companion image column + # (e.g. Lin-Chen/ShareGPT4V, LLaVA-style datasets) + for chat_col in ("conversations", "messages"): + if chat_col not in column_names: + continue + chat_data = sample[chat_col] + if not isinstance(chat_data, list) or len(chat_data) == 0: + continue + first_msg = chat_data[0] + if not isinstance(first_msg, dict): + continue + # Detect ShareGPT (from/value) or ChatML (role/content) keys + msg_text = first_msg.get("value") or first_msg.get("content") + if not isinstance(msg_text, str): + continue + # Check for placeholder anywhere in the conversation + has_image_placeholder = any( + "" in str(m.get("value", "") or m.get("content", "")) + for m in chat_data + if isinstance(m, dict) + ) + if not has_image_placeholder: + continue + # Find companion image column + image_col = None + for col in column_names: + if col == chat_col: + continue + if _keyword_in_column("image", col) or _keyword_in_column("img", col): + image_col = col + break + if image_col: + return { + "format": "sharegpt_with_images", + "needs_conversion": True, + "image_column": image_col, + "text_column": None, + "messages_column": chat_col, + } + # Find image and text columns using metadata filtering # Define metadata patterns to EXCLUDE @@ -584,10 +642,10 @@ def detect_vlm_dataset_structure(dataset): } # Image-related keywords - image_keywords = ['image', 'img', 'photo', 'picture', 'pic', 'visual', 'scan'] + image_keywords = ['image', 'img', 'photo', 'picture', 'pic', 'visual', 'scan', 'file_name', 'filename'] # Text-related keywords - text_keywords = ['text', 'caption', 'description', 'answer', 'output', 'response', 'label'] + text_keywords = ['text', 'caption', 'captions', 'description', 'answer', 'output', 'response', 'label'] def is_metadata_column(col_name): """Check if column name looks like metadata.""" @@ -603,39 +661,92 @@ def detect_vlm_dataset_structure(dataset): return False + def _score_image_candidate(col, sample_value): + """Score a candidate image column by how resolvable its value is.""" + # PIL Image object (highest priority - already loaded) + if hasattr(sample_value, 'size') and hasattr(sample_value, 'mode'): + return 100 + + # Dict with image data (bytes/path from HF Image feature) + if isinstance(sample_value, dict) and ('bytes' in sample_value or 'path' in sample_value): + return 75 + + if isinstance(sample_value, str): + # URL strings + if sample_value.startswith(("http://", "https://")): + return 70 if not is_metadata_column(col) else 55 + # Bare file path + if is_metadata_column(col): + return 30 + return 50 + + return 0 + + def _probe_image_candidate(col, sample_value): + """Quick probe to check if an image candidate is actually reachable. + Returns True if likely valid, False if definitely broken.""" + import os + + # PIL / dict — already loaded, always valid + if not isinstance(sample_value, str): + return True + + # Local file — check it exists + if not sample_value.startswith(("http://", "https://")): + return os.path.exists(sample_value) # bare filenames return False here, that's OK + + # URL — quick HEAD request with short timeout + try: + import urllib.request + req = urllib.request.Request(sample_value, method="HEAD") + resp = urllib.request.urlopen(req, timeout=3) + return resp.status < 400 + except Exception: + return False + def find_image_column(): - """Find image column by filtering out metadata and checking keywords.""" + """Find image column by keyword match + value-based fallback. + When multiple candidates exist, probes them to find one that works.""" candidates = [] + # Pass 1: keyword-matched columns for col in column_names: - col_lower = col.lower() - - # Check if contains image keywords - if any(keyword in col_lower for keyword in image_keywords): - # Verify it actually contains image data + if any(_keyword_in_column(keyword, col) for keyword in image_keywords): sample_value = sample[col] + score = _score_image_candidate(col, sample_value) + if score > 0: + candidates.append((col, score)) - # PIL Image object (highest priority - even if name suggests metadata) - if hasattr(sample_value, 'size') and hasattr(sample_value, 'mode'): - candidates.append((col, 100)) # High priority - actual PIL Image + # Pass 2: value-based fallback — find columns with image URLs/paths + # even if the column name doesn't match image keywords + already = {c[0] for c in candidates} + for col in column_names: + if col in already: + continue + sample_value = sample[col] + if _is_image_value(sample_value): + score = _score_image_candidate(col, sample_value) + # Slightly penalise non-keyword columns so keyword matches win on ties + candidates.append((col, max(score - 5, 1))) - # String (could be path) - but lower priority if name is metadata-like - elif isinstance(sample_value, str): - if is_metadata_column(col): - candidates.append((col, 30)) # Lower priority for metadata names - else: - candidates.append((col, 50)) # Medium priority + if not candidates: + return None - # Dict with image data - elif isinstance(sample_value, dict) and ('bytes' in sample_value or 'path' in sample_value): - candidates.append((col, 75)) # High-medium priority + candidates.sort(key=lambda x: x[1], reverse=True) - # Return highest priority candidate - if candidates: - candidates.sort(key=lambda x: x[1], reverse=True) + # Single candidate or top candidate is PIL/dict — no probing needed + if len(candidates) == 1 or candidates[0][1] >= 75: return candidates[0][0] - return None + # Multiple string-based candidates — probe to find one that actually works + for col, score in candidates: + sample_value = sample[col] + if _probe_image_candidate(col, sample_value): + return col + + # Nothing probed successfully — return highest-scored anyway and let + # conversion handle the error (it may still resolve via hf_hub_download) + return candidates[0][0] def find_text_column(): """Find text column by filtering out metadata and checking keywords.""" @@ -646,10 +757,8 @@ def detect_vlm_dataset_structure(dataset): if is_metadata_column(col): continue - col_lower = col.lower() - - # Check if contains text keywords - if any(keyword in col_lower for keyword in text_keywords): + # Check if contains text keywords (word-boundary match) + if any(_keyword_in_column(keyword, col) for keyword in text_keywords): # Verify it's actually text sample_value = sample[col] @@ -657,6 +766,10 @@ def detect_vlm_dataset_structure(dataset): # Longer text = higher priority (likely content, not just a label) priority = min(len(sample_value), 1000) # Cap at 1000 candidates.append((col, priority)) + elif isinstance(sample_value, list) and len(sample_value) > 0 and isinstance(sample_value[0], str): + # List of strings (e.g. captions list) — lower priority than plain strings + priority = min(len(sample_value[0]), 1000) // 2 + candidates.append((col, priority)) # Return highest priority candidate if candidates: