rustytorch/examples/model_hub_demo.rs

//! Model Hub Demo - One-Command Model Loading
//!
//! This example demonstrates how to load models from various sources:
//! - HuggingFace Hub
//! - RustyTorch Hub
//! - Local directories
//!
//! Run with: cargo run --example model_hub_demo

use rtx_hub::{
    load, load_with_config, LoadConfig, LoadedModel,
    RTXArchitecture, RTXDType, HubResult,
    list_cached_models, cache_size, clear_cache,
    SafeTensors, ShardedSafeTensors,
};
use std::path::PathBuf;

#[tokio::main]
async fn main() -> HubResult<()> {
    // Initialize logging
    tracing_subscriber::fmt::init();

    println!("=== RustyTorch++ Model Hub Demo ===\n");

    // Example 1: Simple one-line loading
    demo_simple_loading().await?;

    // Example 2: Loading with custom configuration
    demo_custom_config().await?;

    // Example 3: Loading from local directory
    demo_local_loading().await?;

    // Example 4: Working with SafeTensors
    demo_safetensors().await?;

    // Example 5: Cache management
    demo_cache_management().await?;

    println!("\n=== Demo Complete ===");
    Ok(())
}

/// Demo 1: Simple one-line model loading
async fn demo_simple_loading() -> HubResult<()> {
    println!("--- Demo 1: Simple Model Loading ---\n");

    // The simplest way to load a model - just provide the model ID
    // This automatically detects if it's a HuggingFace or RustyTorch model
    println!("Loading model: meta-llama/Llama-3.2-1B (simulated)");

    // Note: In production, this would actually download the model
    // For demo purposes, we'll show the API usage

    /*
    let model = load("meta-llama/Llama-3.2-1B").await?;

    println!("Model loaded successfully!");
    println!("  Architecture: {:?}", model.architecture());
    println!("  Hidden size: {}", model.hidden_size());
    println!("  Num layers: {}", model.num_layers());
    println!("  Vocab size: {}", model.vocab_size());
    println!("  Path: {}", model.path.display());
    */

    println!("Example API usage:");
    println!(r#"
    // Load a model with one line
    let model = rtx_hub::load("meta-llama/Llama-3.2-1B").await?;

    // Access model info
    println!("Architecture: {:?}", model.architecture());
    println!("Hidden size: {}", model.hidden_size());
    println!("Num layers: {}", model.num_layers());
    "#);

    Ok(())
}

/// Demo 2: Loading with custom configuration
async fn demo_custom_config() -> HubResult<()> {
    println!("\n--- Demo 2: Custom Configuration ---\n");

    // Build a custom configuration
    let config = LoadConfig::new()
        .with_revision("main")           // Specify git revision
        .with_dtype(RTXDType::BF16)      // Target data type
        .with_device("cuda:0")           // Target device
        .with_cache_dir("/tmp/rtx-models") // Custom cache location
        .with_hf_hub();                   // Force HuggingFace Hub

    println!("Configuration:");
    println!("  Revision: {:?}", config.revision);
    println!("  DType: {:?}", config.dtype);
    println!("  Device: {:?}", config.device);
    println!("  Cache dir: {:?}", config.cache_dir);
    println!("  Use HF Hub: {}", config.use_hf_hub);

    println!("\nExample API usage:");
    println!(r#"
    let config = LoadConfig::new()
        .with_revision("v1.0.0")
        .with_dtype(RTXDType::BF16)
        .with_device("cuda:0")
        .with_auth_token("hf_xxx")  // For private models
        .with_offline(true);        // Use cached only

    let model = rtx_hub::load_with_config("mistralai/Mistral-7B-v0.1", config).await?;
    "#);

    Ok(())
}

/// Demo 3: Loading from local directory
async fn demo_local_loading() -> HubResult<()> {
    println!("\n--- Demo 3: Local Model Loading ---\n");

    println!("Example API usage:");
    println!(r#"
    // Load from a local directory (must contain config.json)
    let model = rtx_hub::load("/path/to/local/model").await?;

    // The model directory should have:
    // - config.json (required)
    // - model.safetensors or model.safetensors.index.json
    // - tokenizer.json (optional)
    // - tokenizer_config.json (optional)
    // - generation_config.json (optional)
    "#);

    Ok(())
}

/// Demo 4: Working with SafeTensors
async fn demo_safetensors() -> HubResult<()> {
    println!("\n--- Demo 4: SafeTensors Support ---\n");

    // Create a simple SafeTensors file in memory
    use rtx_hub::{SafeTensorsBuilder, SafeTensorsDType};

    // Build a SafeTensors file with some test data
    let tensor_data = vec![0u8; 1024]; // 256 floats

    let safetensors_bytes = SafeTensorsBuilder::new()
        .add_tensor("layer.0.weight", SafeTensorsDType::F32, vec![16, 16], tensor_data.clone())
        .add_tensor("layer.0.bias", SafeTensorsDType::F32, vec![16], vec![0u8; 64])
        .with_metadata("format", "pt")
        .with_metadata("framework", "rustytorch")
        .build()?;

    println!("Created SafeTensors file: {} bytes", safetensors_bytes.len());

    // Load it back
    let loaded = SafeTensors::from_bytes(&safetensors_bytes)?;

    println!("Loaded SafeTensors:");
    println!("  Number of tensors: {}", loaded.num_tensors());
    println!("  Tensor names: {:?}", loaded.tensor_names());

    for name in loaded.tensor_names() {
        if let Some(info) = loaded.tensor_info(name) {
            println!("  {} - dtype: {:?}, shape: {:?}", name, info.dtype, info.shape);
        }
    }

    // Validate integrity
    loaded.validate()?;
    println!("  Validation: OK");

    println!("\nExample API usage for loading model weights:");
    println!(r#"
    // After loading a model
    let model = rtx_hub::load("meta-llama/Llama-3.2-1B").await?;

    // Load weights (single file)
    if !model.is_sharded() {
        let weights = model.load_safetensors().await?;
        for name in weights.tensor_names() {
            let (info, data) = weights.tensor_data(name).unwrap();
            println!("{}: {:?} ({} bytes)", name, info.shape, data.len());
        }
    }

    // Load weights (sharded model)
    if model.is_sharded() {
        let weights = model.load_sharded_safetensors().await?;
        println!("Loaded {} shards", weights.num_shards());
    }
    "#);

    Ok(())
}

/// Demo 5: Cache management
async fn demo_cache_management() -> HubResult<()> {
    println!("\n--- Demo 5: Cache Management ---\n");

    // List cached models
    let cached = list_cached_models(None).await?;
    println!("Cached models: {:?}", cached);

    // Get cache size
    let size = cache_size(None).await?;
    println!("Cache size: {} bytes ({:.2} MB)", size, size as f64 / 1_000_000.0);

    println!("\nExample API usage:");
    println!(r#"
    // List all cached models
    let models = rtx_hub::list_cached_models(None).await?;
    for model_id in models {
        println!("Cached: {}", model_id);
    }

    // Get cache size
    let size = rtx_hub::cache_size(None).await?;
    println!("Cache: {:.2} GB", size as f64 / 1_000_000_000.0);

    // Clear cache (be careful!)
    // rtx_hub::clear_cache(None).await?;

    // Use custom cache directory
    let custom_cache = PathBuf::from("/data/models");
    let models = rtx_hub::list_cached_models(Some(custom_cache)).await?;
    "#);

    Ok(())
}

/// Bonus: Architecture detection example
fn demo_architecture_detection() {
    println!("\n--- Bonus: Architecture Detection ---\n");

    let architectures = [
        ("llama", RTXArchitecture::from_hf_model_type("llama")),
        ("mistral", RTXArchitecture::from_hf_model_type("mistral")),
        ("gpt2", RTXArchitecture::from_hf_model_type("gpt2")),
        ("bert", RTXArchitecture::from_hf_model_type("bert")),
        ("phi3", RTXArchitecture::from_hf_model_type("phi3")),
        ("qwen2", RTXArchitecture::from_hf_model_type("qwen2")),
        ("mamba", RTXArchitecture::from_hf_model_type("mamba")),
    ];

    println!("Supported architectures:");
    for (name, arch) in architectures {
        println!("  {} -> {:?}", name, arch);
    }
}