synthethic_multimodal_data/assets.py

from __future__ import annotations

import statistics

from dagster import (
    AssetCheckResult,
    AssetCheckSeverity,
    AssetExecutionContext,
    MaterializeResult,
    asset,
    asset_check,
)
from dagster_hf_datasets import hf_dataset_asset
from datasets import Dataset

# ── Config ─────────────────────────────────────────────────────────────────
#
# SAMPLE_SIZE controls CPU runtime. BLIP-base captioning takes roughly
# 0.5-2s/image on CPU, so 40 images runs in ~1-2 minutes. Increase for
# a richer demo; runtime scales roughly linearly.

SAMPLE_SIZE = 40
NSFW_THRESHOLD = 0.5      # DiffusionDB image_nsfw / prompt_nsfw scores, 0-1
ALIGNMENT_THRESHOLD = 0.2  # MiniLM cosine similarity, prompt vs. generated caption


# ── Step 1: Ingest pre-generated synthetic images ────────────────────────────
#
# DiffusionDB already contains Stable-Diffusion-generated images paired
# with the prompts used to create them. This sidesteps running diffusion
# inference entirely — there's no GPU step in this pipeline. The "synthetic"
# framing comes from the images themselves (AI-generated), not from any
# generation step performed here.

@hf_dataset_asset(
    path="poloclub/diffusiondb",
    config="2m_random_1k",
    split="train",
    group_name="synthetic_multimodal_generation",
    io_manager_key="hf_parquet_io_manager",
)
def diffusiondb_sample(
    context: AssetExecutionContext,
    dataset: Dataset,
) -> MaterializeResult:
    """Ingest a small sample of pre-generated Stable Diffusion images + prompts.

    DiffusionDB's `2m_random_1k` config provides 1,000 random
    (image, prompt, generation-params) triples. Only the first
    `SAMPLE_SIZE` rows are kept to bound CPU runtime for downstream
    captioning steps.
    """
    n = min(SAMPLE_SIZE, len(dataset))
    sample = dataset.select(range(n))

    context.log.info("Loaded DiffusionDB sample: %s / %s rows", n, len(dataset))
    context.log.info("Columns: %s", sample.column_names)

    return MaterializeResult(
        value=sample,
        metadata={
            "rows": n,
            "source_rows_available": len(dataset),
            "columns": sample.column_names,
            "source_dataset": "poloclub/diffusiondb",
            "config": "2m_random_1k",
            "fingerprint": sample._fingerprint,
        },
    )


# ── Step 2: NSFW-based quality filtering ──────────────────────────────────────
#
# DiffusionDB ships precomputed image_nsfw / prompt_nsfw scores from the
# original collection pipeline. Filtering on these requires no model —
# it's free quality filtering using existing metadata.

@asset(
    group_name="synthetic_multimodal_generation",
    io_manager_key="hf_parquet_io_manager",
)
def nsfw_filtered(
    context: AssetExecutionContext,
    diffusiondb_sample: Dataset,
) -> MaterializeResult:
    """Drop rows flagged by DiffusionDB's precomputed NSFW scores.

    `image_nsfw` and `prompt_nsfw` are floats in [0, 1] (or -1 if not
    scored). Rows scoring >= NSFW_THRESHOLD on either field are dropped.
    Unscored (-1) rows pass through.
    """
    before = len(diffusiondb_sample)

    def is_safe(example: dict) -> bool:
        img_nsfw = example.get("image_nsfw", -1.0)
        prompt_nsfw = example.get("prompt_nsfw", -1.0)
        if img_nsfw is not None and img_nsfw >= NSFW_THRESHOLD:
            return False
        if prompt_nsfw is not None and prompt_nsfw >= NSFW_THRESHOLD:
            return False
        return True

    filtered = diffusiondb_sample.filter(is_safe, desc="NSFW filter")
    after = len(filtered)

    context.log.info("NSFW filter: %s → %s rows (%s dropped)", before, after, before - after)
    context.add_output_metadata(
        {
            "rows_in": before,
            "rows_out": after,
            "dropped": before - after,
            "nsfw_threshold": NSFW_THRESHOLD,
        }
    )
    return MaterializeResult(
        value=filtered,
        metadata={
            "rows": after,
            "rows_in": before,
            "rows_out": after,
            "dropped": before - after,
            "nsfw_threshold": NSFW_THRESHOLD,
        },
    )


# ── Step 3: Synthetic caption generation (BLIP, CPU) ──────────────────────────
#
# This is the "VLM evaluation" step from the original spec, inverted:
# rather than using a VLM to *score* generated images, we use a small
# captioning VLM (BLIP-base, ~990MB) to *generate* a caption for each
# image. BLIP-base runs comfortably on CPU for small batches.

_blip_processor = None
_blip_model = None


def _load_blip():
    """Lazily load BLIP captioning model (downloads ~990MB on first run)."""
    global _blip_processor, _blip_model
    if _blip_model is None:
        from transformers import BlipForConditionalGeneration, BlipProcessor

        _blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
        _blip_model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        _blip_model.eval()
    return _blip_processor, _blip_model


@asset(
    group_name="synthetic_multimodal_generation",
    io_manager_key="hf_parquet_io_manager",
)
def synthetic_captions(
    context: AssetExecutionContext,
    nsfw_filtered: Dataset,
) -> MaterializeResult:
    """Generate a caption for each image using BLIP-base (CPU inference).

    Adds a `generated_caption` column. Each caption is generated
    independently — no batching — to keep memory bounded on CPU.
    """
    import torch

    processor, model = _load_blip()
    captions: list[str] = []

    for i, example in enumerate(nsfw_filtered):
        img = example["image"].convert("RGB")
        inputs = processor(img, return_tensors="pt")
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=30)
        caption = processor.decode(out[0], skip_special_tokens=True)
        captions.append(caption)

        if i % 10 == 0:
            context.log.info("Captioned %s / %s images", i, len(nsfw_filtered))

    captioned = nsfw_filtered.add_column("generated_caption", captions)

    avg_len = statistics.mean(len(c.split()) for c in captions)
    context.log.info("Captioning complete. Avg caption length: %.1f words", avg_len)
    context.add_output_metadata(
        {
            "rows": len(captioned),
            "model": "Salesforce/blip-image-captioning-base",
            "avg_caption_length_words": round(avg_len, 2),
        }
    )
    return MaterializeResult(
        value=captioned,
        metadata={
            "rows": len(captioned),
            "model": "Salesforce/blip-image-captioning-base",
            "avg_caption_length_words": round(avg_len, 2),
        },
    )


# ── Step 4: Caption-prompt alignment scoring (MiniLM, CPU) ────────────────────
#
# Automated evaluation step: scores how semantically aligned the
# BLIP-generated caption is with the original DiffusionDB prompt
# using sentence-embedding cosine similarity. all-MiniLM-L6-v2 is
# ~90MB and runs fast on CPU.

_sentence_model = None


def _load_sentence_model():
    """Lazily load the sentence-embedding model (downloads ~90MB on first run)."""
    global _sentence_model
    if _sentence_model is None:
        from sentence_transformers import SentenceTransformer

        _sentence_model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
    return _sentence_model


@asset(
    group_name="synthetic_multimodal_generation",
    io_manager_key="hf_parquet_io_manager",
)
def caption_alignment_scores(
    context: AssetExecutionContext,
    synthetic_captions: Dataset,
) -> MaterializeResult:
    """Score semantic alignment between original prompt and generated caption.

    Adds an `alignment_score` column: cosine similarity between
    sentence embeddings of `prompt` and `generated_caption`, in [-1, 1].
    Low scores indicate the BLIP caption diverged significantly from
    what the prompt asked for — useful for flagging generation failures
    or captioning failures.
    """
    from sentence_transformers import util

    model = _load_sentence_model()
    scores: list[float] = []

    for i, example in enumerate(synthetic_captions):
        emb = model.encode(
            [example["prompt"], example["generated_caption"]],
            convert_to_tensor=True,
        )
        score = float(util.cos_sim(emb[0], emb[1]))
        scores.append(score)

        if i % 10 == 0:
            context.log.info("Scored %s / %s pairs", i, len(synthetic_captions))

    scored = synthetic_captions.add_column("alignment_score", scores)

    context.log.info(
        "Alignment scores — min: %.3f, max: %.3f, mean: %.3f",
        min(scores), max(scores), statistics.mean(scores),
    )
    context.add_output_metadata(
        {
            "rows": len(scored),
            "model": "all-MiniLM-L6-v2",
            "alignment_score_min": round(min(scores), 3),
            "alignment_score_max": round(max(scores), 3),
            "alignment_score_mean": round(statistics.mean(scores), 3),
        }
    )
    return MaterializeResult(
        value=scored,
        metadata={
            "rows": len(scored),
            "model": "all-MiniLM-L6-v2",
            "alignment_score_min": round(min(scores), 3),
            "alignment_score_max": round(max(scores), 3),
            "alignment_score_mean": round(statistics.mean(scores), 3),
        },
    )


# ── Step 5: Final filtered synthetic dataset ──────────────────────────────────

@asset(
    group_name="synthetic_multimodal_generation",
    io_manager_key="hf_parquet_io_manager",
)
def synthetic_dataset_final(
    context: AssetExecutionContext,
    caption_alignment_scores: Dataset,
) -> MaterializeResult:
    """Filter to rows whose generated caption aligns well with the original prompt.

    Rows with `alignment_score < ALIGNMENT_THRESHOLD` are dropped. The
    output is a synthetic image-caption-prompt triple dataset suitable
    for downstream fine-tuning or evaluation harness use.
    """
    before = len(caption_alignment_scores)

    filtered = caption_alignment_scores.filter(
        lambda ex: ex["alignment_score"] >= ALIGNMENT_THRESHOLD,
        desc="Alignment filter",
    )
    after = len(filtered)

    keep_columns = ["image", "prompt", "generated_caption", "alignment_score", "seed", "cfg", "sampler"]
    available = [c for c in keep_columns if c in filtered.column_names]
    drop_columns = [c for c in filtered.column_names if c not in available]
    final = filtered.remove_columns(drop_columns) if drop_columns else filtered

    context.log.info(
        "Alignment filter: %s → %s rows (%s dropped, threshold=%.2f)",
        before, after, before - after, ALIGNMENT_THRESHOLD,
    )
    context.add_output_metadata(
        {
            "rows_in": before,
            "rows_out": after,
            "dropped": before - after,
            "alignment_threshold": ALIGNMENT_THRESHOLD,
            "columns": final.column_names,
        }
    )
    return MaterializeResult(
        value=final,
        metadata={
            "rows": len(final),
            "rows_in": before,
            "rows_out": after,
            "dropped": before - after,
            "alignment_threshold": ALIGNMENT_THRESHOLD,
            "columns": final.column_names,
        },
    )


# ── Step 6: Generation report ─────────────────────────────────────────────────

@asset(
    group_name="synthetic_multimodal_generation",
)
def synthetic_generation_report(
    context: AssetExecutionContext,
    diffusiondb_sample: Dataset,
    nsfw_filtered: Dataset,
    caption_alignment_scores: Dataset,
    synthetic_dataset_final: Dataset,
) -> MaterializeResult:
    """Funnel report across the synthetic generation + filtering pipeline."""
    stages = {
        "raw_sample": len(diffusiondb_sample),
        "after_nsfw_filter": len(nsfw_filtered),
        "after_captioning": len(caption_alignment_scores),
        "after_alignment_filter": len(synthetic_dataset_final),
    }

    scores = [ex["alignment_score"] for ex in caption_alignment_scores]

    report = {
        "stages": stages,
        "final_retention_pct": round(stages["after_alignment_filter"] / stages["raw_sample"] * 100, 2)
        if stages["raw_sample"] else 0.0,
        "alignment_score_mean": round(statistics.mean(scores), 3) if scores else None,
        "alignment_score_min": round(min(scores), 3) if scores else None,
        "alignment_score_max": round(max(scores), 3) if scores else None,
    }

    context.log.info("Synthetic generation funnel: %s", stages)
    context.add_output_metadata(
        {
            **{f"stage_{k}": v for k, v in stages.items()},
            "final_retention_pct": report["final_retention_pct"],
            "alignment_score_mean": report["alignment_score_mean"],
        }
    )
    return MaterializeResult(
        value=report,
        metadata={
            **{f"stage_{k}": v for k, v in stages.items()},
            "final_retention_pct": report["final_retention_pct"],
            "alignment_score_mean": report["alignment_score_mean"],
        },
    )


# ── Asset checks ──────────────────────────────────────────────────────────────

@asset_check(
    asset=synthetic_dataset_final,
    description="Final dataset has non-empty generated captions for every row",
)
def check_no_empty_captions(synthetic_dataset_final: Dataset) -> AssetCheckResult:
    empty = sum(
        1 for ex in synthetic_dataset_final
        if not ex.get("generated_caption") or len(ex["generated_caption"].strip()) == 0
    )
    return AssetCheckResult(
        passed=empty == 0,
        severity=AssetCheckSeverity.ERROR,
        metadata={"empty_caption_count": empty},
    )


@asset_check(
    asset=synthetic_dataset_final,
    description="Mean alignment score in final dataset is above the filter threshold",
)
def check_mean_alignment(synthetic_dataset_final: Dataset) -> AssetCheckResult:
    if len(synthetic_dataset_final) == 0:
        return AssetCheckResult(
            passed=False,
            severity=AssetCheckSeverity.WARN,
            metadata={"reason": "no rows survived alignment filter"},
        )

    scores = [ex["alignment_score"] for ex in synthetic_dataset_final]
    mean_score = statistics.mean(scores)

    return AssetCheckResult(
        passed=mean_score >= ALIGNMENT_THRESHOLD,
        severity=AssetCheckSeverity.WARN,
        metadata={
            "mean_alignment_score": round(mean_score, 3),
            "threshold": ALIGNMENT_THRESHOLD,
            "rows": len(synthetic_dataset_final),
        },
    )