from __future__ import annotations import json import statistics from collections import Counter from pathlib import Path from dagster import AssetExecutionContext, MaterializeResult, asset from dagster_hf_datasets import hf_dataset_asset from datasets import Dataset # ── Step 1: Ingest ──────────────────────────────────────────────────────────── @hf_dataset_asset( path="nlphuji/flickr30k", split="test", group_name="multimodal_profiling", io_manager_key="hf_parquet_io_manager", ) def flickr30k_raw( context: AssetExecutionContext, dataset: Dataset, ) -> MaterializeResult: """Ingest Flickr30k from the Hub. Flickr30k contains 31,783 images each paired with 5 captions, stored as PIL Image objects in the `image` column and lists in the `caption` column. """ context.log.info("Loaded Flickr30k: %s rows", len(dataset)) context.log.info("Columns: %s", dataset.column_names) context.log.info("Features: %s", dataset.features) return MaterializeResult( value=dataset, metadata={ "rows": len(dataset), "columns": dataset.column_names, "source_dataset": "nlphuji/flickr30k", "split": "test", "fingerprint": dataset._fingerprint, }, ) # ── Step 2: Image statistics ────────────────────────────────────────────────── @asset( group_name="multimodal_profiling", io_manager_key="hf_parquet_io_manager", ) def image_stats( context: AssetExecutionContext, flickr30k_raw: Dataset, ) -> MaterializeResult: """Extract per-image resolution and aspect ratio statistics. Iterates over PIL Image objects in the dataset, recording width, height, aspect ratio, and mode for each image. Results are returned as a flat Dataset for downstream use and stored via the IO manager. """ records = [] for i, example in enumerate(flickr30k_raw): img = example["image"] # PIL.Image width, height = img.size aspect = round(width / height, 4) if height > 0 else 0.0 records.append( { "idx": i, "width": width, "height": height, "aspect_ratio": aspect, "mode": img.mode, "megapixels": round((width * height) / 1_000_000, 4), } ) if i % 1000 == 0: context.log.info("Processed %s / %s images", i, len(flickr30k_raw)) widths = [r["width"] for r in records] heights = [r["height"] for r in records] aspects = [r["aspect_ratio"] for r in records] mode_counts = Counter(r["mode"] for r in records) context.log.info( "Width — min: %s, max: %s, mean: %.1f", min(widths), max(widths), statistics.mean(widths), ) context.log.info( "Height — min: %s, max: %s, mean: %.1f", min(heights), max(heights), statistics.mean(heights), ) context.log.info("Color modes: %s", dict(mode_counts)) stats_dataset = Dataset.from_list(records) context.add_output_metadata( { "image_count": len(records), "width_min": min(widths), "width_max": max(widths), "width_mean": round(statistics.mean(widths), 1), "height_min": min(heights), "height_max": max(heights), "height_mean": round(statistics.mean(heights), 1), "aspect_ratio_mean": round(statistics.mean(aspects), 4), "color_modes": str(dict(mode_counts)), } ) return MaterializeResult( value=stats_dataset, metadata={ "rows": len(stats_dataset), "image_count": len(records), "width_min": min(widths), "width_max": max(widths), "width_mean": round(statistics.mean(widths), 1), "height_min": min(heights), "height_max": max(heights), "height_mean": round(statistics.mean(heights), 1), "aspect_ratio_mean": round(statistics.mean(aspects), 4), "color_modes": str(dict(mode_counts)), }, ) # ── Step 3: Caption statistics ──────────────────────────────────────────────── @asset( group_name="multimodal_profiling", io_manager_key="hf_parquet_io_manager", ) def caption_stats( context: AssetExecutionContext, flickr30k_raw: Dataset, ) -> MaterializeResult: """Compute per-example caption length and vocabulary statistics. Each Flickr30k example has a list of 5 captions. This asset flattens them and computes token counts, unique word counts, and average caption length per example. """ records = [] all_tokens: list[str] = [] for i, example in enumerate(flickr30k_raw): captions: list[str] = example["caption"] # list of 5 strings token_counts = [len(c.split()) for c in captions] avg_tokens = statistics.mean(token_counts) all_tokens.extend(tok for c in captions for tok in c.lower().split()) records.append( { "idx": i, "num_captions": len(captions), "avg_caption_length": round(avg_tokens, 2), "min_caption_length": min(token_counts), "max_caption_length": max(token_counts), } ) vocab_size = len(set(all_tokens)) all_lengths = [r["avg_caption_length"] for r in records] context.log.info("Total vocabulary size: %s unique tokens", vocab_size) context.log.info( "Caption length — min: %.1f, max: %.1f, mean: %.1f", min(all_lengths), max(all_lengths), statistics.mean(all_lengths), ) stats_dataset = Dataset.from_list(records) context.add_output_metadata( { "examples_analyzed": len(records), "vocabulary_size": vocab_size, "avg_caption_length_mean": round(statistics.mean(all_lengths), 2), "avg_caption_length_min": round(min(all_lengths), 2), "avg_caption_length_max": round(max(all_lengths), 2), "total_captions": sum(r["num_captions"] for r in records), } ) return MaterializeResult( value=stats_dataset, metadata={ "rows": len(stats_dataset), "examples_analyzed": len(records), "vocabulary_size": vocab_size, "avg_caption_length_mean": round(statistics.mean(all_lengths), 2), "avg_caption_length_min": round(min(all_lengths), 2), "avg_caption_length_max": round(max(all_lengths), 2), "total_captions": sum(r["num_captions"] for r in records), }, ) # ── Step 4: Thumbnail gallery ───────────────────────────────────────────────── @asset( group_name="multimodal_profiling", ) def sample_gallery( context: AssetExecutionContext, flickr30k_raw: Dataset, ) -> MaterializeResult: """Save a thumbnail gallery of 16 sample images to disk. Writes 128×128 JPEG thumbnails to `.dagster_hf_storage/sample_gallery/`. Returns a manifest dict with file paths and the first caption for each. """ gallery_dir = Path(".dagster_hf_storage/sample_gallery") gallery_dir.mkdir(parents=True, exist_ok=True) sample_indices = list(range(0, min(16, len(flickr30k_raw)))) manifest = [] for i in sample_indices: example = flickr30k_raw[i] img = example["image"].copy() img.thumbnail((128, 128)) out_path = gallery_dir / f"sample_{i:04d}.jpg" img.save(out_path, format="JPEG", quality=85) manifest.append( { "idx": i, "path": str(out_path), "caption": example["caption"][0] if example["caption"] else "", "original_size": example["image"].size, "thumbnail_size": img.size, } ) manifest_path = gallery_dir / "manifest.json" manifest_path.write_text(json.dumps(manifest, indent=2)) context.log.info("Saved %s thumbnails to %s", len(manifest), gallery_dir) return MaterializeResult( value={"thumbnails": manifest, "gallery_dir": str(gallery_dir)}, metadata={ "thumbnail_count": len(manifest), "gallery_dir": str(gallery_dir), "manifest_path": str(manifest_path), }, ) # ── Step 5: Health report ───────────────────────────────────────────────────── @asset( group_name="multimodal_profiling", ) def dataset_health_report( context: AssetExecutionContext, flickr30k_raw: Dataset, image_stats: Dataset, caption_stats: Dataset, ) -> MaterializeResult: """Combine image and caption statistics into a single health report. Flags potential quality issues: - Images with unusual aspect ratios (< 0.2 or > 5.0) - Captions shorter than 4 tokens on average - Missing or empty captions """ # Aspect ratio outliers extreme_aspect = [ row for row in image_stats if row["aspect_ratio"] < 0.2 or row["aspect_ratio"] > 5.0 ] # Short captions short_captions = [ row for row in caption_stats if row["avg_caption_length"] < 4.0 ] # Missing captions missing_captions = sum( 1 for ex in flickr30k_raw if not ex.get("caption") or len(ex["caption"]) == 0 ) total = len(flickr30k_raw) report = { "total_examples": total, "extreme_aspect_ratio_count": len(extreme_aspect), "extreme_aspect_ratio_pct": round(len(extreme_aspect) / total * 100, 2), "short_caption_count": len(short_captions), "short_caption_pct": round(len(short_captions) / total * 100, 2), "missing_caption_count": missing_captions, "health_score": round( 100 - (len(extreme_aspect) / total * 30) - (len(short_captions) / total * 40) - (missing_captions / total * 30), 1, ), } context.log.info("Dataset health report: %s", report) context.add_output_metadata(report) return MaterializeResult( value=report, metadata=report, ) # ── Step 6: LLaVA-150K (Modern Instruction-Tuned Vision-Language) ──────────── @hf_dataset_asset( path="liuhaotian/llava-instruct-150k", split="train", group_name="multimodal_profiling", io_manager_key="hf_parquet_io_manager", ) def llava_instruct_raw( context: AssetExecutionContext, dataset: Dataset, ) -> MaterializeResult: """Ingest LLaVA-150K instruction-tuning dataset. LLaVA-150K contains 150K image-instruction-response pairs, designed for instruction-tuning vision-language models. Unlike raw image-caption pairs (Flickr30K), this data has explicit instruction-following structure (Q&A format). Modern alternative to generic image-caption datasets for VLM training. """ context.log.info("Loaded LLaVA-150K: %s rows", len(dataset)) context.log.info("Columns: %s", dataset.column_names) return MaterializeResult( value=dataset.select(range(min(5000, len(dataset)))), metadata={ "rows": min(5000, len(dataset)), "columns": dataset.column_names, "source_dataset": "liuhaotian/llava-instruct-150k", "split": "train", "description": "Instruction-tuning data for vision-language models", }, ) @asset( group_name="multimodal_profiling", io_manager_key="hf_parquet_io_manager", ) def llava_instruction_stats( context: AssetExecutionContext, llava_instruct_raw: Dataset, ) -> MaterializeResult: """Extract instruction and response statistics from LLaVA data. Unlike Flickr30K (multiple captions per image), LLaVA has instruction-response pairs. This asset computes: - Instruction complexity (token count, question type) - Response length distribution - Instruction-response alignment """ records = [] for i, example in enumerate(llava_instruct_raw): # LLaVA structure: 'image', 'conversations' (list of turns) conversations = example.get("conversations", []) # Typically alternates user (instruction) and assistant (response) user_turn = None assistant_turn = None for turn in conversations: role = turn.get("from") text = turn.get("value", "").strip() if role == "human": user_turn = text elif role == "gpt": assistant_turn = text if user_turn and assistant_turn: records.append( { "idx": i, "instruction_tokens": len(user_turn.split()), "response_tokens": len(assistant_turn.split()), "instruction_length": len(user_turn), "response_length": len(assistant_turn), "is_question": user_turn.rstrip().endswith("?"), } ) if i % 1000 == 0: context.log.info("Processed %s / %s examples", i, len(llava_instruct_raw)) instr_tokens = [r["instruction_tokens"] for r in records if records] resp_tokens = [r["response_tokens"] for r in records if records] context.log.info( "Instruction length — min: %s, max: %s, mean: %.1f", min(instr_tokens) if instr_tokens else 0, max(instr_tokens) if instr_tokens else 0, statistics.mean(instr_tokens) if instr_tokens else 0, ) context.log.info( "Response length — min: %s, max: %s, mean: %.1f", min(resp_tokens) if resp_tokens else 0, max(resp_tokens) if resp_tokens else 0, statistics.mean(resp_tokens) if resp_tokens else 0, ) stats_dataset = Dataset.from_list(records) context.add_output_metadata( { "example_count": len(records), "instruction_tokens_mean": round(statistics.mean(instr_tokens), 1) if instr_tokens else 0, "response_tokens_mean": round(statistics.mean(resp_tokens), 1) if resp_tokens else 0, "question_fraction": round(sum(1 for r in records if r["is_question"]) / len(records), 2) if records else 0, } ) return MaterializeResult( value=stats_dataset, metadata={ "rows": len(stats_dataset), "example_count": len(records), "instruction_tokens_mean": round(statistics.mean(instr_tokens), 1) if instr_tokens else 0, "response_tokens_mean": round(statistics.mean(resp_tokens), 1) if resp_tokens else 0, "question_fraction": round( sum(1 for r in records if r["is_question"]) / len(records), 2, ) if records else 0, }, ) @asset(group_name="multimodal_profiling") def llava_quality_profile( context: AssetExecutionContext, llava_instruct_raw: Dataset, llava_instruction_stats: Dataset, ) -> MaterializeResult: """Profile quality metrics specific to instruction-tuned VLM data. Checks: - Instruction-response pair validity - Length balance (not too short, not excessively long) - Response quality indicators """ total = len(llava_instruct_raw) stats_records = list(llava_instruction_stats) # Check response quality very_short_responses = sum(1 for r in stats_records if r["response_tokens"] < 5) very_long_responses = sum(1 for r in stats_records if r["response_tokens"] > 500) balanced_responses = total - very_short_responses - very_long_responses # Instruction diversity questions = sum(1 for r in stats_records if r["is_question"]) question_pct = round(questions / len(stats_records) * 100, 1) if stats_records else 0 profile = { "total_examples": total, "valid_instruction_response_pairs": len(stats_records), "very_short_responses_count": very_short_responses, "very_long_responses_count": very_long_responses, "balanced_responses": balanced_responses, "balanced_response_pct": round(balanced_responses / total * 100, 2), "question_percentage": question_pct, "instruction_complexity_score": round( statistics.mean([r["instruction_tokens"] for r in stats_records]) / 10, 1 ) if stats_records else 0, } context.log.info("LLaVA quality profile: %s", profile) return MaterializeResult( value=profile, metadata=profile, )