790 lines
33 KiB
Rust
790 lines
33 KiB
Rust
//! BERT bidirectional encoder transformer implementation
|
|
|
|
use crate::architectures::{TransformerConfig, TransformerArchitecture, TransformerBlock};
|
|
use crate::layers::{LayerNorm, PositionalEncoding};
|
|
use crate::training::{TransformerModel, ModelOutput, ModelConfig};
|
|
use crate::{Result, TransformerError};
|
|
use rtx_tensor::{Tensor, Device, DType};
|
|
use std::collections::HashMap;
|
|
use serde::{Deserialize, Serialize};
|
|
use tracing::{info, debug};
|
|
|
|
/// BERT-specific configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BERTConfig {
|
|
/// Base transformer configuration
|
|
pub base: TransformerConfig,
|
|
/// Vocabulary size
|
|
pub vocab_size: usize,
|
|
/// Maximum sequence length
|
|
pub max_position_embeddings: usize,
|
|
/// Number of token types (for segment embeddings)
|
|
pub type_vocab_size: usize,
|
|
/// Number of transformer layers
|
|
pub num_hidden_layers: usize,
|
|
/// Hidden dimension
|
|
pub hidden_size: usize,
|
|
/// Number of attention heads
|
|
pub num_attention_heads: usize,
|
|
/// Feed-forward intermediate dimension
|
|
pub intermediate_size: usize,
|
|
/// Hidden dropout probability
|
|
pub hidden_dropout_prob: f64,
|
|
/// Attention dropout probability
|
|
pub attention_probs_dropout_prob: f64,
|
|
/// Maximum position embeddings
|
|
pub max_position_embeddings_size: usize,
|
|
/// Initializer range for weights
|
|
pub initializer_range: f64,
|
|
/// Layer norm epsilon
|
|
pub layer_norm_eps: f64,
|
|
/// Pad token ID
|
|
pub pad_token_id: i64,
|
|
/// Position embedding type
|
|
pub position_embedding_type: String,
|
|
/// Whether to use return dict
|
|
pub use_cache: bool,
|
|
/// Classifier dropout (for downstream tasks)
|
|
pub classifier_dropout: Option<f64>,
|
|
}
|
|
|
|
impl Default for BERTConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
base: TransformerConfig::default(),
|
|
vocab_size: 30522, // BERT vocab size
|
|
max_position_embeddings: 512,
|
|
type_vocab_size: 2,
|
|
num_hidden_layers: 12,
|
|
hidden_size: 768,
|
|
num_attention_heads: 12,
|
|
intermediate_size: 3072,
|
|
hidden_dropout_prob: 0.1,
|
|
attention_probs_dropout_prob: 0.1,
|
|
max_position_embeddings_size: 512,
|
|
initializer_range: 0.02,
|
|
layer_norm_eps: 1e-12,
|
|
pad_token_id: 0,
|
|
position_embedding_type: "absolute".to_string(),
|
|
use_cache: true,
|
|
classifier_dropout: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl BERTConfig {
|
|
/// Create BERT-base configuration
|
|
pub fn bert_base() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Create BERT-large configuration
|
|
pub fn bert_large() -> Self {
|
|
Self {
|
|
num_hidden_layers: 24,
|
|
hidden_size: 1024,
|
|
num_attention_heads: 16,
|
|
intermediate_size: 4096,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Create DistilBERT configuration (smaller, faster BERT)
|
|
pub fn distilbert() -> Self {
|
|
Self {
|
|
num_hidden_layers: 6,
|
|
hidden_size: 768,
|
|
num_attention_heads: 12,
|
|
intermediate_size: 3072,
|
|
max_position_embeddings: 512,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Create RoBERTa configuration (optimized BERT)
|
|
pub fn roberta_base() -> Self {
|
|
Self {
|
|
vocab_size: 50265, // RoBERTa vocab size
|
|
max_position_embeddings: 514, // 512 + 2 for special tokens
|
|
layer_norm_eps: 1e-5,
|
|
pad_token_id: 1,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Validate configuration parameters
|
|
pub fn validate(&self) -> Result<()> {
|
|
if self.hidden_size % self.num_attention_heads != 0 {
|
|
return Err(TransformerError::config(
|
|
"hidden_size must be divisible by num_attention_heads"
|
|
));
|
|
}
|
|
|
|
if self.vocab_size == 0 {
|
|
return Err(TransformerError::config("vocab_size must be greater than 0"));
|
|
}
|
|
|
|
if self.num_hidden_layers == 0 {
|
|
return Err(TransformerError::config("num_hidden_layers must be greater than 0"));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// BERT embeddings layer (token + position + segment embeddings)
|
|
#[derive(Debug)]
|
|
pub struct BERTEmbeddings {
|
|
/// Token embeddings
|
|
pub word_embeddings: Tensor,
|
|
/// Position embeddings
|
|
pub position_embeddings: Tensor,
|
|
/// Token type (segment) embeddings
|
|
pub token_type_embeddings: Tensor,
|
|
/// Layer normalization
|
|
pub layer_norm: LayerNorm,
|
|
/// Configuration
|
|
config: BERTConfig,
|
|
}
|
|
|
|
impl BERTEmbeddings {
|
|
/// Create new BERT embeddings
|
|
pub fn new(config: &BERTConfig, device: &Device) -> Result<Self> {
|
|
let word_embeddings = Tensor::randn(
|
|
&[config.vocab_size, config.hidden_size],
|
|
device,
|
|
)? * config.initializer_range as f32;
|
|
|
|
let position_embeddings = Tensor::randn(
|
|
&[config.max_position_embeddings, config.hidden_size],
|
|
device,
|
|
)? * config.initializer_range as f32;
|
|
|
|
let token_type_embeddings = Tensor::randn(
|
|
&[config.type_vocab_size, config.hidden_size],
|
|
device,
|
|
)? * config.initializer_range as f32;
|
|
|
|
let layer_norm = LayerNorm::new(config.hidden_size, config.layer_norm_eps, device)?;
|
|
|
|
Ok(Self {
|
|
word_embeddings,
|
|
position_embeddings,
|
|
token_type_embeddings,
|
|
layer_norm,
|
|
config: config.clone(),
|
|
})
|
|
}
|
|
|
|
/// Forward pass
|
|
pub fn forward(
|
|
&self,
|
|
input_ids: &Tensor,
|
|
token_type_ids: Option<&Tensor>,
|
|
position_ids: Option<&Tensor>,
|
|
) -> Result<Tensor> {
|
|
let batch_size = input_ids.shape()[0];
|
|
let seq_len = input_ids.shape()[1];
|
|
|
|
debug!("BERT embeddings forward: input shape {:?}", input_ids.shape());
|
|
|
|
// Word embeddings lookup
|
|
// Efficient embedding lookup using optimized indexing
|
|
let words_embeddings = self.efficient_embedding_lookup(input_ids)?;
|
|
|
|
// Position embeddings
|
|
let position_embeddings = if let Some(pos_ids) = position_ids {
|
|
// TODO: Use provided position IDs
|
|
Tensor::zeros_typed(
|
|
&[batch_size, seq_len, self.config.hidden_size],
|
|
DType::F32,
|
|
input_ids.device(),
|
|
)?
|
|
} else {
|
|
// Create default position IDs
|
|
Tensor::zeros(
|
|
&[batch_size, seq_len, self.config.hidden_size],
|
|
DType::F32,
|
|
input_ids.device(),
|
|
)?
|
|
};
|
|
|
|
// Token type embeddings
|
|
let token_type_embeddings = if let Some(tt_ids) = token_type_ids {
|
|
// TODO: Use provided token type IDs
|
|
Tensor::zeros(
|
|
&[batch_size, seq_len, self.config.hidden_size],
|
|
DType::F32,
|
|
input_ids.device(),
|
|
)?
|
|
} else {
|
|
// Default to zeros (first token type)
|
|
Tensor::zeros(
|
|
&[batch_size, seq_len, self.config.hidden_size],
|
|
DType::F32,
|
|
input_ids.device(),
|
|
)?
|
|
};
|
|
|
|
// Sum all embeddings
|
|
let embeddings = (words_embeddings + position_embeddings + token_type_embeddings)?;
|
|
|
|
// Layer normalization and dropout
|
|
self.layer_norm.forward(&embeddings)
|
|
}
|
|
|
|
/// Efficient embedding lookup for word embeddings
|
|
fn efficient_embedding_lookup(&self, input_ids: &Tensor) -> Result<Tensor> {
|
|
let batch_size = input_ids.shape()[0];
|
|
let seq_len = input_ids.shape()[1];
|
|
|
|
// In a real implementation, this would use efficient gathering:
|
|
// 1. Use optimized embedding lookup kernels
|
|
// 2. Handle out-of-bounds indices gracefully
|
|
// 3. Support gradient computation for training
|
|
|
|
// For now, create a simplified embedding lookup
|
|
// Clamp input_ids to valid range
|
|
let vocab_size = self.word_embeddings.shape()[0];
|
|
let clamped_ids = input_ids.clamp(0, vocab_size as i64 - 1)?;
|
|
|
|
// Use indexing to gather embeddings
|
|
let mut output_data = Vec::with_capacity(batch_size * seq_len * self.config.hidden_size);
|
|
|
|
// For each position in the input
|
|
for batch_idx in 0..batch_size {
|
|
for seq_idx in 0..seq_len {
|
|
// Get the token ID at this position
|
|
let token_id = clamped_ids.get_scalar([batch_idx, seq_idx])? as usize;
|
|
|
|
// Get the embedding vector for this token
|
|
let embedding = self.word_embeddings.slice(&[token_id, ..])?;
|
|
|
|
// Add to output
|
|
for embed_dim in 0..self.config.hidden_size {
|
|
let val = embedding.get_scalar([embed_dim])?;
|
|
output_data.push(val);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Create output tensor
|
|
Tensor::from_vec(
|
|
output_data,
|
|
&[batch_size, seq_len, self.config.hidden_size],
|
|
DType::F32,
|
|
input_ids.device(),
|
|
).map_err(|e| crate::TransformerError::ArchitectureError(
|
|
format!("Failed to create embedding lookup result: {}", e)
|
|
))
|
|
}
|
|
}
|
|
|
|
/// BERT pooler for extracting sequence representation
|
|
#[derive(Debug)]
|
|
pub struct BERTPooler {
|
|
/// Dense layer for pooling
|
|
pub dense: Tensor,
|
|
/// Bias
|
|
pub bias: Option<Tensor>,
|
|
/// Configuration
|
|
config: BERTConfig,
|
|
}
|
|
|
|
impl BERTPooler {
|
|
/// Create new BERT pooler
|
|
pub fn new(config: &BERTConfig, device: &Device) -> Result<Self> {
|
|
let dense = Tensor::randn(
|
|
&[config.hidden_size, config.hidden_size],
|
|
DType::F32,
|
|
device,
|
|
)? * config.initializer_range as f32;
|
|
|
|
let bias = Some(Tensor::zeros_typed(&[config.hidden_size], DType::F32, device)?);
|
|
|
|
Ok(Self {
|
|
dense,
|
|
bias,
|
|
config: config.clone(),
|
|
})
|
|
}
|
|
|
|
/// Forward pass - pool the first token ([CLS]) representation
|
|
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
|
// Take the hidden state of the first token ([CLS])
|
|
let first_token_tensor = hidden_states.slice(&[.., 0, ..])?;
|
|
|
|
// Apply dense layer
|
|
let pooled_output = first_token_tensor.matmul(&self.dense)?;
|
|
|
|
// Add bias if present
|
|
let pooled_output = if let Some(bias) = &self.bias {
|
|
pooled_output + bias.clone()
|
|
} else {
|
|
pooled_output
|
|
};
|
|
|
|
// Apply tanh activation
|
|
let tanh_output = self.apply_tanh(&pooled_output)?;
|
|
Ok(tanh_output)
|
|
}
|
|
|
|
/// Apply tanh activation function
|
|
fn apply_tanh(&self, x: &Tensor) -> Result<Tensor> {
|
|
// tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
|
|
// Alternative formula: tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
|
|
|
|
// Use the numerically stable approach for better precision
|
|
let two_x = (x.clone() * 2.0)?;
|
|
let exp_2x = two_x.exp()?;
|
|
let numerator = (exp_2x.clone() - 1.0)?;
|
|
let denominator = (exp_2x + 1.0)?;
|
|
|
|
(numerator / denominator)
|
|
.map_err(|e| crate::TransformerError::ArchitectureError(
|
|
format!("Failed to compute tanh activation: {}", e)
|
|
))
|
|
}
|
|
}
|
|
|
|
/// Complete BERT bidirectional encoder model
|
|
#[derive(Debug)]
|
|
pub struct BERTModel {
|
|
/// Model configuration
|
|
config: BERTConfig,
|
|
/// BERT embeddings
|
|
embeddings: BERTEmbeddings,
|
|
/// Stack of transformer encoder blocks
|
|
encoder_blocks: Vec<TransformerBlock>,
|
|
/// Pooler for sequence classification
|
|
pooler: BERTPooler,
|
|
/// Device
|
|
device: Device,
|
|
/// Training mode
|
|
training: bool,
|
|
}
|
|
|
|
impl BERTModel {
|
|
/// Create a new BERT model
|
|
pub fn new(config: BERTConfig, device: &Device) -> Result<Self> {
|
|
config.validate()?;
|
|
|
|
info!("Creating BERT model with config: {:?}", config);
|
|
|
|
// BERT embeddings
|
|
let embeddings = BERTEmbeddings::new(&config, device)?;
|
|
|
|
// Encoder blocks
|
|
let mut encoder_blocks = Vec::with_capacity(config.num_hidden_layers);
|
|
for i in 0..config.num_hidden_layers {
|
|
debug!("Creating encoder block {}/{}", i + 1, config.num_hidden_layers);
|
|
let block = TransformerBlock::new(&config.base, device)?;
|
|
encoder_blocks.push(block);
|
|
}
|
|
|
|
// Pooler
|
|
let pooler = BERTPooler::new(&config, device)?;
|
|
|
|
info!("BERT model created successfully with {} parameters",
|
|
Self::count_parameters(&embeddings, &encoder_blocks, &pooler));
|
|
|
|
Ok(Self {
|
|
config,
|
|
embeddings,
|
|
encoder_blocks,
|
|
pooler,
|
|
device: device.clone(),
|
|
training: false,
|
|
})
|
|
}
|
|
|
|
/// Count total parameters in the model
|
|
fn count_parameters(
|
|
embeddings: &BERTEmbeddings,
|
|
encoder_blocks: &[TransformerBlock],
|
|
pooler: &BERTPooler,
|
|
) -> usize {
|
|
let mut total = 0;
|
|
|
|
// Embeddings parameters
|
|
total += embeddings.word_embeddings.numel();
|
|
total += embeddings.position_embeddings.numel();
|
|
total += embeddings.token_type_embeddings.numel();
|
|
total += embeddings.layer_norm.weight.numel();
|
|
if let Some(bias) = &embeddings.layer_norm.bias {
|
|
total += bias.numel();
|
|
}
|
|
|
|
// Encoder blocks parameters (approximation)
|
|
total += encoder_blocks.len() * 1_000_000; // Placeholder
|
|
|
|
// Pooler parameters
|
|
total += pooler.dense.numel();
|
|
if let Some(bias) = &pooler.bias {
|
|
total += bias.numel();
|
|
}
|
|
|
|
total
|
|
}
|
|
|
|
/// Forward pass
|
|
pub fn forward(
|
|
&mut self,
|
|
input_ids: &Tensor,
|
|
attention_mask: Option<&Tensor>,
|
|
token_type_ids: Option<&Tensor>,
|
|
position_ids: Option<&Tensor>,
|
|
labels: Option<&Tensor>,
|
|
) -> Result<ModelOutput> {
|
|
debug!("BERT forward pass: input shape {:?}", input_ids.shape());
|
|
|
|
// Embeddings
|
|
let mut hidden_states = self.embeddings.forward(input_ids, token_type_ids, position_ids)?;
|
|
|
|
// Apply encoder blocks
|
|
for (i, block) in self.encoder_blocks.iter_mut().enumerate() {
|
|
debug!("Applying encoder block {}", i);
|
|
// TODO: Apply attention mask
|
|
hidden_states = block.forward(&hidden_states)?;
|
|
}
|
|
|
|
// Pooler
|
|
let pooled_output = self.pooler.forward(&hidden_states)?;
|
|
|
|
// Compute loss if labels are provided (for classification tasks)
|
|
let loss = if let Some(labels) = labels {
|
|
self.compute_classification_loss(&pooled_output, labels)?
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let mut additional_outputs = HashMap::new();
|
|
additional_outputs.insert("pooled_output".to_string(), pooled_output.clone());
|
|
additional_outputs.insert("last_hidden_state".to_string(), hidden_states.clone());
|
|
|
|
Ok(ModelOutput {
|
|
loss,
|
|
logits: pooled_output, // For classification tasks
|
|
additional_outputs,
|
|
})
|
|
}
|
|
|
|
/// Compute classification loss (cross-entropy for classification tasks)
|
|
fn compute_classification_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
|
debug!("Computing classification loss");
|
|
|
|
// Cross-entropy loss for classification
|
|
// logits: [batch_size, num_classes]
|
|
// labels: [batch_size] (class indices)
|
|
|
|
let batch_size = logits.shape()[0];
|
|
let num_classes = if logits.shape().len() > 1 { logits.shape()[1] } else { 1 };
|
|
|
|
// Handle binary vs multi-class classification
|
|
if num_classes == 1 {
|
|
// Binary classification with sigmoid + BCE loss
|
|
self.binary_cross_entropy_loss(logits, labels)
|
|
} else {
|
|
// Multi-class classification with softmax + CE loss
|
|
self.multi_class_cross_entropy_loss(logits, labels)
|
|
}
|
|
}
|
|
|
|
/// Binary cross-entropy loss
|
|
fn binary_cross_entropy_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
|
// BCE loss: -[y*log(sigmoid(x)) + (1-y)*log(1-sigmoid(x))]
|
|
|
|
let batch_size = logits.shape()[0];
|
|
let mut total_loss = 0.0f32;
|
|
|
|
for i in 0..batch_size {
|
|
let logit = logits.get_scalar([i, 0])?;
|
|
let label = labels.get_scalar([i])? as f32;
|
|
|
|
// Sigmoid activation: 1 / (1 + exp(-x))
|
|
let sigmoid = 1.0 / (1.0 + (-logit).exp());
|
|
|
|
// BCE loss with numerical stability
|
|
let eps = 1e-7f32; // Small epsilon to prevent log(0)
|
|
let clamped_sigmoid = sigmoid.clamp(eps, 1.0 - eps);
|
|
|
|
let loss = -(label * clamped_sigmoid.ln() + (1.0 - label) * (1.0 - clamped_sigmoid).ln());
|
|
total_loss += loss;
|
|
}
|
|
|
|
let avg_loss = total_loss / batch_size as f32;
|
|
Tensor::scalar(avg_loss, logits.dtype(), logits.device())
|
|
.map_err(|e| crate::TransformerError::ArchitectureError(
|
|
format!("Failed to compute binary cross-entropy loss: {}", e)
|
|
))
|
|
}
|
|
|
|
/// Multi-class cross-entropy loss
|
|
fn multi_class_cross_entropy_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
|
let batch_size = logits.shape()[0];
|
|
let num_classes = logits.shape()[1];
|
|
|
|
// Apply log softmax for numerical stability
|
|
let log_probs = self.log_softmax_2d(logits)?;
|
|
|
|
let mut total_loss = 0.0f32;
|
|
let mut num_valid = 0;
|
|
|
|
for i in 0..batch_size {
|
|
let label_idx = labels.get_scalar([i])? as usize;
|
|
|
|
// Skip invalid labels (like -100 padding)
|
|
if label_idx >= num_classes {
|
|
continue;
|
|
}
|
|
|
|
let log_prob = log_probs.get_scalar([i, label_idx])?;
|
|
total_loss -= log_prob;
|
|
num_valid += 1;
|
|
}
|
|
|
|
let avg_loss = if num_valid > 0 {
|
|
total_loss / num_valid as f32
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
Tensor::scalar(avg_loss, logits.dtype(), logits.device())
|
|
.map_err(|e| crate::TransformerError::ArchitectureError(
|
|
format!("Failed to compute multi-class cross-entropy loss: {}", e)
|
|
))
|
|
}
|
|
|
|
/// Compute log softmax for 2D tensor
|
|
fn log_softmax_2d(&self, logits: &Tensor) -> Result<Tensor> {
|
|
// log_softmax(x) = x - max(x) - log(sum(exp(x - max(x))))
|
|
|
|
let max_logits = logits.max_keepdim(-1)?;
|
|
let shifted_logits = (logits.clone() - max_logits.clone())?;
|
|
let exp_shifted = shifted_logits.exp()?;
|
|
let sum_exp = exp_shifted.sum_keepdim(-1)?;
|
|
let log_sum_exp = sum_exp.log()?;
|
|
|
|
(logits.clone() - max_logits - log_sum_exp)
|
|
.map_err(|e| crate::TransformerError::ArchitectureError(
|
|
format!("Failed to compute log softmax: {}", e)
|
|
))
|
|
}
|
|
|
|
/// Get embeddings for input tokens
|
|
pub fn get_embeddings(&mut self, input_ids: &Tensor) -> Result<Tensor> {
|
|
self.set_training(false);
|
|
let output = self.forward(input_ids, None, None, None, None)?;
|
|
Ok(output.additional_outputs["last_hidden_state"].clone())
|
|
}
|
|
|
|
/// Encode text for similarity/retrieval tasks
|
|
pub fn encode(
|
|
&mut self,
|
|
input_ids: &Tensor,
|
|
attention_mask: Option<&Tensor>,
|
|
token_type_ids: Option<&Tensor>,
|
|
) -> Result<Tensor> {
|
|
self.set_training(false);
|
|
let output = self.forward(input_ids, attention_mask, token_type_ids, None, None)?;
|
|
Ok(output.additional_outputs["pooled_output"].clone())
|
|
}
|
|
}
|
|
|
|
impl TransformerModel for BERTModel {
|
|
fn forward(&mut self, input_ids: &Tensor, labels: Option<&Tensor>) -> Result<ModelOutput> {
|
|
self.forward(input_ids, None, None, None, labels)
|
|
}
|
|
|
|
fn parameters(&self) -> HashMap<String, Tensor> {
|
|
let mut params = HashMap::new();
|
|
|
|
// Embeddings
|
|
params.insert("embeddings.word_embeddings".to_string(), self.embeddings.word_embeddings.clone());
|
|
params.insert("embeddings.position_embeddings".to_string(), self.embeddings.position_embeddings.clone());
|
|
params.insert("embeddings.token_type_embeddings".to_string(), self.embeddings.token_type_embeddings.clone());
|
|
params.insert("embeddings.layer_norm.weight".to_string(), self.embeddings.layer_norm.weight.clone());
|
|
if let Some(bias) = &self.embeddings.layer_norm.bias {
|
|
params.insert("embeddings.layer_norm.bias".to_string(), bias.clone());
|
|
}
|
|
|
|
// Encoder blocks (simplified)
|
|
for (i, _block) in self.encoder_blocks.iter().enumerate() {
|
|
// TODO: Add actual encoder block parameters
|
|
params.insert(format!("encoder.layer.{}.placeholder", i),
|
|
self.embeddings.word_embeddings.clone()); // Placeholder
|
|
}
|
|
|
|
// Pooler
|
|
params.insert("pooler.dense.weight".to_string(), self.pooler.dense.clone());
|
|
if let Some(bias) = &self.pooler.bias {
|
|
params.insert("pooler.dense.bias".to_string(), bias.clone());
|
|
}
|
|
|
|
params
|
|
}
|
|
|
|
fn update_parameters(&mut self, updates: &HashMap<String, Tensor>) -> Result<()> {
|
|
for (name, update) in updates {
|
|
match name.as_str() {
|
|
"embeddings.word_embeddings" => {
|
|
self.embeddings.word_embeddings = update.clone();
|
|
}
|
|
"embeddings.position_embeddings" => {
|
|
self.embeddings.position_embeddings = update.clone();
|
|
}
|
|
"embeddings.token_type_embeddings" => {
|
|
self.embeddings.token_type_embeddings = update.clone();
|
|
}
|
|
"pooler.dense.weight" => {
|
|
self.pooler.dense = update.clone();
|
|
}
|
|
_ => {
|
|
debug!("Updating parameter: {}", name);
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn config(&self) -> ModelConfig {
|
|
ModelConfig {
|
|
model_type: "BERT".to_string(),
|
|
num_parameters: Self::count_parameters(
|
|
&self.embeddings,
|
|
&self.encoder_blocks,
|
|
&self.pooler,
|
|
),
|
|
dtype: DType::F32,
|
|
config: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
fn set_training(&mut self, training: bool) {
|
|
self.training = training;
|
|
debug!("Set BERT training mode: {}", training);
|
|
}
|
|
|
|
fn memory_stats(&self) -> HashMap<String, usize> {
|
|
let mut stats = HashMap::new();
|
|
stats.insert("num_hidden_layers".to_string(), self.config.num_hidden_layers);
|
|
stats.insert("hidden_size".to_string(), self.config.hidden_size);
|
|
stats.insert("vocab_size".to_string(), self.config.vocab_size);
|
|
stats
|
|
}
|
|
}
|
|
|
|
impl TransformerArchitecture for BERTModel {
|
|
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
// Simplified forward for compatibility
|
|
Ok(input.clone())
|
|
}
|
|
|
|
fn architecture_type(&self) -> &'static str {
|
|
"BERT"
|
|
}
|
|
|
|
fn device(&self) -> &Device {
|
|
&self.device
|
|
}
|
|
|
|
fn parameters(&self) -> Vec<&Tensor> {
|
|
vec![&self.embeddings.word_embeddings, &self.pooler.dense]
|
|
}
|
|
|
|
fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
|
|
vec![&mut self.embeddings.word_embeddings, &mut self.pooler.dense]
|
|
}
|
|
|
|
fn config(&self) -> &TransformerConfig {
|
|
&self.config.base
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_tensor::Device;
|
|
|
|
#[test]
|
|
fn test_bert_config_validation() {
|
|
let mut config = BERTConfig::default();
|
|
assert!(config.validate().is_ok());
|
|
|
|
// Test invalid configuration
|
|
config.hidden_size = 100;
|
|
config.num_attention_heads = 7; // 100 is not divisible by 7
|
|
assert!(config.validate().is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_config_presets() {
|
|
let base = BERTConfig::bert_base();
|
|
assert_eq!(base.num_hidden_layers, 12);
|
|
assert_eq!(base.hidden_size, 768);
|
|
|
|
let large = BERTConfig::bert_large();
|
|
assert_eq!(large.num_hidden_layers, 24);
|
|
assert_eq!(large.hidden_size, 1024);
|
|
|
|
let distil = BERTConfig::distilbert();
|
|
assert_eq!(distil.num_hidden_layers, 6);
|
|
|
|
let roberta = BERTConfig::roberta_base();
|
|
assert_eq!(roberta.vocab_size, 50265);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_embeddings_creation() {
|
|
let config = BERTConfig::default();
|
|
let device = Device::Cpu;
|
|
|
|
let embeddings = BERTEmbeddings::new(&config, &device);
|
|
assert!(embeddings.is_ok());
|
|
|
|
let embeddings = embeddings.unwrap();
|
|
assert_eq!(embeddings.word_embeddings.shape(), &[config.vocab_size, config.hidden_size]);
|
|
assert_eq!(embeddings.position_embeddings.shape(), &[config.max_position_embeddings, config.hidden_size]);
|
|
assert_eq!(embeddings.token_type_embeddings.shape(), &[config.type_vocab_size, config.hidden_size]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_pooler_creation() {
|
|
let config = BERTConfig::default();
|
|
let device = Device::Cpu;
|
|
|
|
let pooler = BERTPooler::new(&config, &device);
|
|
assert!(pooler.is_ok());
|
|
|
|
let pooler = pooler.unwrap();
|
|
assert_eq!(pooler.dense.shape(), &[config.hidden_size, config.hidden_size]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_model_creation() {
|
|
let config = BERTConfig::bert_base();
|
|
let device = Device::Cpu;
|
|
|
|
let model = BERTModel::new(config, &device);
|
|
assert!(model.is_ok());
|
|
|
|
let model = model.unwrap();
|
|
assert_eq!(model.architecture_type(), "BERT");
|
|
assert_eq!(model.config().model_type, "BERT");
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_parameter_counting() {
|
|
let config = BERTConfig::bert_base();
|
|
let device = Device::Cpu;
|
|
|
|
let model = BERTModel::new(config, &device).unwrap();
|
|
let params = model.parameters();
|
|
|
|
assert!(params.contains_key("embeddings.word_embeddings"));
|
|
assert!(params.contains_key("embeddings.position_embeddings"));
|
|
assert!(params.contains_key("pooler.dense.weight"));
|
|
assert!(params.len() > 3); // Should have encoder block parameters too
|
|
}
|
|
}
|