Large Scale Datasets

import asyncio
import time
from contextlib import contextmanager
from rich.console import Console

from kura.cache import DiskCacheStrategy
from kura.types import Conversation
from kura.summarisation import SummaryModel, summarise_conversations
from kura.k_means import MiniBatchKmeansClusteringMethod
from kura.cluster import (
    ClusterDescriptionModel,
    generate_base_clusters_from_conversation_summaries,
)
from kura.meta_cluster import MetaClusterModel, reduce_clusters_from_base_clusters
from kura.dimensionality import HDBUMAP, reduce_dimensionality_from_clusters
from kura.checkpoints import (
    HFDatasetCheckpointManager,
    JSONLCheckpointManager,
    ParquetCheckpointManager,
)
from kura.checkpoints import MultiCheckpointManager


class TimerManager:
    """Track timing for each pipeline step."""

    def __init__(self):
        self.timings = {}

    @contextmanager
    def timer(self, message):
        start_time = time.time()
        yield
        duration = time.time() - start_time
        self.timings[message] = duration
        print(f"{message} took {duration:.2f} seconds")

    def print_summary(self):
        print(f"\n{'=' * 60}")
        print(f"{'TIMING SUMMARY':^60}")
        print(f"{'=' * 60}")
        total_time = sum(self.timings.values())

        for operation, duration in self.timings.items():
            percentage = (duration / total_time * 100) if total_time > 0 else 0
            print(f"{operation:<40} {duration:>8.2f}s ({percentage:>5.1f}%)")

        print(f"{'-' * 60}")
        print(f"{'Total Time':<40} {total_time:>8.2f}s")
        print(f"{'=' * 60}\n")


async def main():
    console = Console()
    timer_manager = TimerManager()
    CHECKPOINT_DIR = "./large_scale_checkpoints"

    # Step 1: Configure models for high throughput
    summary_model = SummaryModel(
        console=console,
        max_concurrent_requests=100,  # Increased for speed
        cache=DiskCacheStrategy(cache_dir="./.summary_cache"),
    )

    # Use MiniBatch K-means for large-scale clustering
    minibatch_kmeans = MiniBatchKmeansClusteringMethod(
        clusters_per_group=10,  # Target items per cluster
        batch_size=1000,  # Process 1000 embeddings at a time
        max_iter=100,  # Iterations for convergence
        random_state=42,
    )

    cluster_model = ClusterDescriptionModel(
        console=console,
        clustering_method=minibatch_kmeans,
    )

    meta_cluster_model = MetaClusterModel(
        console=console,
        max_concurrent_requests=100,
    )

    dimensionality_model = HDBUMAP()

    # Step 2: Set up multi-format checkpointing
    # Saves to all formats simultaneously for redundancy and performance
    hf_manager = HFDatasetCheckpointManager(f"{CHECKPOINT_DIR}/hf", enabled=True)
    parquet_manager = ParquetCheckpointManager(
        f"{CHECKPOINT_DIR}/parquet", enabled=True
    )
    jsonl_manager = JSONLCheckpointManager(f"{CHECKPOINT_DIR}/jsonl", enabled=True)

    checkpoint_manager = MultiCheckpointManager(
        [hf_manager, parquet_manager, jsonl_manager],
        save_strategy="all_enabled",  # Save to all formats
        load_strategy="first_found",  # Load from fastest available
    )

    print(f"Using MultiCheckpointManager with {len(checkpoint_manager.managers)} backends:")
    for manager in checkpoint_manager.managers:
        print(f"  - {type(manager).__name__}: {manager.checkpoint_dir}")
    print()

    # Step 3: Load large dataset
    with timer_manager.timer("Loading conversations"):
        conversations = Conversation.from_hf_dataset(
            "ivanleomk/synthetic-gemini-conversations",
            split="train",
            max_conversations=10000,  # Scale to 10k conversations
        )

    console.print(f"[bold green]✓ Loaded {len(conversations)} conversations[/bold green]\n")

    # Step 4: Process pipeline with checkpoints
    console.print("[bold blue]Step 1: Generating summaries...[/bold blue]")
    with timer_manager.timer("Summarization"):
        summaries = await summarise_conversations(
            conversations,
            model=summary_model,
            checkpoint_manager=checkpoint_manager,
        )
    console.print(f"[bold green]✓ Generated {len(summaries)} summaries[/bold green]\n")

    console.print("[bold blue]Step 2: Clustering conversations...[/bold blue]")
    with timer_manager.timer("Clustering"):
        clusters = await generate_base_clusters_from_conversation_summaries(
            summaries,
            model=cluster_model,
            checkpoint_manager=checkpoint_manager,
        )
    console.print(f"[bold green]✓ Created {len(clusters)} clusters[/bold green]\n")

    console.print("[bold blue]Step 3: Meta clustering...[/bold blue]")
    with timer_manager.timer("Meta clustering"):
        reduced_clusters = await reduce_clusters_from_base_clusters(
            clusters,
            model=meta_cluster_model,
            checkpoint_manager=checkpoint_manager,
        )
    console.print(f"[bold green]✓ Reduced to {len(reduced_clusters)} meta clusters[/bold green]\n")

    console.print("[bold blue]Step 4: Dimensionality reduction...[/bold blue]")
    with timer_manager.timer("Dimensionality reduction"):
        projected_clusters = await reduce_dimensionality_from_clusters(
            reduced_clusters,
            model=dimensionality_model,
            checkpoint_manager=checkpoint_manager,
        )
    console.print(
        f"[bold green]✓ Projected {len(projected_clusters)} clusters[/bold green]\n"
    )

    # Step 5: Show statistics
    timer_manager.print_summary()

    # Show checkpoint statistics
    print("\nCheckpoint Statistics:")
    stats = checkpoint_manager.get_stats()
    for mgr_stat in stats["managers"]:
        print(f"  - {mgr_stat['type']}: {mgr_stat.get('checkpoint_count', 'N/A')} files")

    return summaries, clusters, reduced_clusters, projected_clusters


if __name__ == "__main__":
    results = asyncio.run(main())