Struct BertModel

pub struct BertModel<F: Float + Debug + ScalarOperand + Send + Sync> { /* private fields */ }

BERT model implementation

Implementations

impl<F: Float + Debug + ScalarOperand + Send + Sync> BertModel<F>

pub fn new(config: BertConfig) -> Result<Self>

Create a new BERT model

Examples found in repository

examples/bert_example.rs (line 18)

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()
    );
    println!(
        "BERT-Base hidden dimension: {}",
        base_pooled_output.shape()[1]
    );

    println!("\nBERT example completed successfully!");

    Ok(())
}

pub fn bert_base_uncased() -> Result<Self>

Create a BERT-Base-Uncased model

Examples found in repository

examples/bert_example.rs (line 42)

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()
    );
    println!(
        "BERT-Base hidden dimension: {}",
        base_pooled_output.shape()[1]
    );

    println!("\nBERT example completed successfully!");

    Ok(())
}

pub fn bert_large_uncased() -> Result<Self>

Create a BERT-Large-Uncased model

pub fn custom( vocab_size: usize, hidden_size: usize, num_hidden_layers: usize, num_attention_heads: usize, ) -> Result<Self>

Create a custom BERT model

pub fn get_sequence_output( &self, input: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Get sequence output (last layer hidden states)

pub fn get_pooled_output( &self, input: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Get pooled output (for classification tasks)

Examples found in repository

examples/bert_example.rs (line 35)

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()
    );
    println!(
        "BERT-Base hidden dimension: {}",
        base_pooled_output.shape()[1]
    );

    println!("\nBERT example completed successfully!");

    Ok(())
}

Trait Implementations

impl<F: Float + Debug + ScalarOperand + Send + Sync> Layer<F> for BertModel<F>

fn forward(&self, input: &Array<F, IxDyn>) -> Result<Array<F, IxDyn>>

Forward pass of the layer

fn backward( &self, _input: &Array<F, IxDyn>, grad_output: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Backward pass of the layer to compute gradients

fn update(&mut self, learning_rate: F) -> Result<()>

Update the layer parameters with the given gradients

fn as_any(&self) -> &dyn Any

Get the layer as a dyn Any for downcasting

fn as_any_mut(&mut self) -> &mut dyn Any

Get the layer as a mutable dyn Any for downcasting

fn params(&self) -> Vec<Array<F, IxDyn>>

Get the parameters of the layer

fn gradients(&self) -> Vec<Array<F, IxDyn>>

Get the gradients of the layer parameters

fn set_gradients(&mut self, _gradients: &[Array<F, IxDyn>]) -> Result<()>

Set the gradients of the layer parameters

fn set_params(&mut self, _params: &[Array<F, IxDyn>]) -> Result<()>

Set the parameters of the layer

fn set_training(&mut self, _training: bool)

Set the layer to training mode (true) or evaluation mode (false)

fn is_training(&self) -> bool

Get the current training mode

fn layer_type(&self) -> &str

Get the type of the layer (e.g., “Dense”, “Conv2D”)

fn parameter_count(&self) -> usize

Get the number of trainable parameters in this layer

fn layer_description(&self) -> String

Get a detailed description of this layer