scirs2-neural 0.3.2

use ndarray::{Array, IxDyn};
use scirs2_neural::layers::Layer;
use scirs2_neural::models::{BertConfig, BertModel};

#[allow(dead_code)]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("BERT Model Example");
    // Create a small BERT model for demonstration
    println!("Creating a small BERT model...");
    let config = BertConfig::custom(
        10000, // vocab_size
        128,   // hidden_size
        2,     // num_hidden_layers
        2,     // num_attention_heads
    );
    let model = BertModel::<f32>::new(config)?;
    // Create dummy input (batch_size=2, seq_len=16)
    // Input tensor contains token IDs
    let input = Array::from_shape_fn(
        IxDyn(&[2, 16]),
        |_| rand::random::<f32>() * 100.0, // Random token IDs between 0 and 100
    println!("Input shape: {:?}", input.shape());
    // Get sequence output (hidden states)
    let sequence_output = model.forward(&input)?;
    println!("Sequence output shape: {:?}", sequence_output.shape());
    // Get pooled output (for classification tasks)
    let pooled_output = model.get_pooled_output(&input)?;
    println!("Pooled output shape: {:?}", pooled_output.shape());
    // Let's create a BERT-Base model
    println!("\nCreating a BERT-Base model...");
    let bert_base = BertModel::<f32>::bert_base_uncased()?;
    // Create dummy input for a longer sequence
    let base_input = Array::from_shape_fn(
        IxDyn(&[1, 64]),
        |_| rand::random::<f32>() * 1000.0, // Random token IDs
    println!("BERT-Base input shape: {:?}", base_input.shape());
    // Forward pass to get pooled output
    let base_pooled_output = bert_base.get_pooled_output(&base_input)?;
    println!(
        "BERT-Base pooled output shape: {:?}",
        base_pooled_output.shape()
        "BERT-Base hidden dimension: {}",
        base_pooled_output.shape()[1]
    println!("\nBERT example completed successfully!");
    Ok(())
}