Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,491 @@
|
||||
//! BERT task-specific heads
|
||||
//!
|
||||
//! This module implements various task-specific heads for BERT models:
|
||||
//! - Masked Language Modeling (MLM) head
|
||||
//! - Sequence classification head
|
||||
//! - Token classification head (NER, POS tagging)
|
||||
//! - Question answering head (SQuAD-style)
|
||||
|
||||
use crate::{Result, TransformerError};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
|
||||
use super::config::BertConfig;
|
||||
use super::feedforward::BertIntermediate;
|
||||
use super::layer_norm::BertLayerNorm;
|
||||
|
||||
/// BERT masked language modeling head
|
||||
#[derive(Debug)]
|
||||
pub struct BertMLMHead {
|
||||
pub hidden_size: usize,
|
||||
pub vocab_size: usize,
|
||||
/// Dense layer for hidden state transformation: [`hidden_size`, `hidden_size`]
|
||||
pub dense: Option<Tensor>,
|
||||
/// Dense layer bias: [`hidden_size`]
|
||||
pub dense_bias: Option<Tensor>,
|
||||
/// Layer normalization for predictions
|
||||
pub layer_norm: Option<BertLayerNorm>,
|
||||
/// Decoder weights for vocabulary projection: [`hidden_size`, `vocab_size`]
|
||||
pub decoder: Option<Tensor>,
|
||||
/// Decoder bias: [`vocab_size`]
|
||||
pub decoder_bias: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl BertMLMHead {
|
||||
/// Create new MLM head
|
||||
pub fn new(config: &BertConfig, device: &Device) -> Result<Self> {
|
||||
// Dense layer for hidden state transformation
|
||||
let init_std = (2.0 / (config.hidden_size + config.hidden_size) as f32).sqrt();
|
||||
let dense = Tensor::randn(&[config.hidden_size, config.hidden_size], device)?
|
||||
.mul_scalar(init_std)?;
|
||||
let dense_bias = Tensor::zeros([config.hidden_size], device)?;
|
||||
|
||||
// Layer normalization
|
||||
let layer_norm = BertLayerNorm::new(config.hidden_size, config.layer_norm_eps, device)?;
|
||||
|
||||
// Decoder (vocabulary projection) layer
|
||||
let decoder_init_std = (2.0 / (config.hidden_size + config.vocab_size) as f32).sqrt();
|
||||
let decoder = Tensor::randn(&[config.hidden_size, config.vocab_size], device)?
|
||||
.mul_scalar(decoder_init_std)?;
|
||||
let decoder_bias = Tensor::zeros([config.vocab_size], device)?;
|
||||
|
||||
Ok(Self {
|
||||
hidden_size: config.hidden_size,
|
||||
vocab_size: config.vocab_size,
|
||||
dense: Some(dense),
|
||||
dense_bias: Some(dense_bias),
|
||||
layer_norm: Some(layer_norm),
|
||||
decoder: Some(decoder),
|
||||
decoder_bias: Some(decoder_bias),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of MLM head
|
||||
pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
|
||||
let dense_weight = self
|
||||
.dense
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("MLM dense layer not initialized"))?;
|
||||
let dense_bias = self
|
||||
.dense_bias
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("MLM dense bias not initialized"))?;
|
||||
let layer_norm = self
|
||||
.layer_norm
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("MLM layer norm not initialized"))?;
|
||||
let decoder_weight = self
|
||||
.decoder
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("MLM decoder not initialized"))?;
|
||||
let decoder_bias = self
|
||||
.decoder_bias
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("MLM decoder bias not initialized"))?;
|
||||
|
||||
// Transform hidden states: hidden_states @ dense + dense_bias
|
||||
let transformed = hidden_states
|
||||
.matmul(dense_weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
let transformed = transformed
|
||||
.add(dense_bias)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
|
||||
// Apply GELU activation
|
||||
let activated = BertIntermediate::gelu_activation(&transformed)?;
|
||||
|
||||
// Apply layer normalization
|
||||
let normalized = layer_norm.forward(&activated)?;
|
||||
|
||||
// Project to vocabulary: normalized @ decoder + decoder_bias
|
||||
let logits = normalized
|
||||
.matmul(decoder_weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
let logits = logits
|
||||
.add(decoder_bias)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
|
||||
Ok(logits)
|
||||
}
|
||||
|
||||
/// Compute masked language modeling loss (cross-entropy)
|
||||
pub fn compute_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
||||
// logits shape: [batch_size, seq_len, vocab_size]
|
||||
// labels shape: [batch_size, seq_len]
|
||||
|
||||
let batch_size = logits.shape().dims()[0];
|
||||
let seq_len = logits.shape().dims()[1];
|
||||
let vocab_size = logits.shape().dims()[2];
|
||||
|
||||
// Flatten logits and labels for cross-entropy computation
|
||||
let logits_flat = logits.reshape([batch_size * seq_len, vocab_size])?;
|
||||
let labels_flat = labels.reshape([batch_size * seq_len])?;
|
||||
|
||||
// Compute softmax probabilities
|
||||
let logits_data = logits_flat.to_cpu()?;
|
||||
let labels_data = labels_flat.to_cpu()?;
|
||||
|
||||
let mut total_loss = 0.0;
|
||||
let mut valid_tokens = 0;
|
||||
|
||||
// Compute cross-entropy loss for each position
|
||||
for i in 0..(batch_size * seq_len) {
|
||||
let label = labels_data[i] as i32;
|
||||
|
||||
// Skip positions with -100 label (mask token)
|
||||
if label == -100 {
|
||||
continue;
|
||||
}
|
||||
|
||||
valid_tokens += 1;
|
||||
|
||||
// Extract logits for this position
|
||||
let start_idx = i * vocab_size;
|
||||
let position_logits: Vec<f32> = logits_data[start_idx..start_idx + vocab_size].to_vec();
|
||||
|
||||
// Compute softmax for numerical stability
|
||||
let max_logit = position_logits
|
||||
.iter()
|
||||
.fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_logits: Vec<f32> = position_logits
|
||||
.iter()
|
||||
.map(|&x| (x - max_logit).exp())
|
||||
.collect();
|
||||
let sum_exp: f32 = exp_logits.iter().sum();
|
||||
|
||||
// Cross-entropy loss: -log(p_true)
|
||||
if label >= 0 && (label as usize) < vocab_size {
|
||||
let prob = exp_logits[label as usize] / sum_exp;
|
||||
total_loss += -prob.ln();
|
||||
}
|
||||
}
|
||||
|
||||
// Average loss over valid tokens
|
||||
let average_loss = if valid_tokens > 0 {
|
||||
total_loss / valid_tokens as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
Tensor::scalar(average_loss, rtx_tensor::DType::F32, logits.device())
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// BERT sequence classification head
|
||||
#[derive(Debug)]
|
||||
pub struct BertClassificationHead {
|
||||
pub hidden_size: usize,
|
||||
pub num_labels: usize,
|
||||
pub dropout_prob: f64,
|
||||
/// Classifier dense layer: [`hidden_size`, `num_labels`]
|
||||
pub classifier: Option<Tensor>,
|
||||
/// Classifier bias: [`num_labels`]
|
||||
pub bias: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl BertClassificationHead {
|
||||
/// Create new classification head
|
||||
pub fn new(
|
||||
hidden_size: usize,
|
||||
num_labels: usize,
|
||||
dropout_prob: f64,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
// Initialize classifier layer with Xavier initialization
|
||||
let init_std = (2.0 / (hidden_size + num_labels) as f32).sqrt();
|
||||
let classifier = Tensor::randn(&[hidden_size, num_labels], device)?.mul_scalar(init_std)?;
|
||||
let bias = Tensor::zeros([num_labels], device)?;
|
||||
|
||||
Ok(Self {
|
||||
hidden_size,
|
||||
num_labels,
|
||||
dropout_prob,
|
||||
classifier: Some(classifier),
|
||||
bias: Some(bias),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of classification head
|
||||
pub fn forward(&self, pooled_output: &Tensor) -> Result<Tensor> {
|
||||
let classifier_weight = self
|
||||
.classifier
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("Classifier not initialized"))?;
|
||||
let bias = self
|
||||
.bias
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("Classifier bias not initialized"))?;
|
||||
|
||||
// Apply dropout (simplified - in practice would check training mode)
|
||||
// For now, skip dropout during inference
|
||||
|
||||
// Apply classifier: pooled_output @ classifier + bias
|
||||
let logits = pooled_output
|
||||
.matmul(classifier_weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
let logits = logits
|
||||
.add(bias)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
|
||||
Ok(logits)
|
||||
}
|
||||
|
||||
/// Compute cross-entropy loss for classification
|
||||
pub fn compute_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
||||
// logits shape: [batch_size, num_labels]
|
||||
// labels shape: [batch_size]
|
||||
|
||||
let batch_size = logits.shape().dims()[0];
|
||||
let num_labels = logits.shape().dims()[1];
|
||||
|
||||
let logits_data = logits.to_cpu()?;
|
||||
let labels_data = labels.to_cpu()?;
|
||||
|
||||
let mut total_loss = 0.0;
|
||||
|
||||
// Compute cross-entropy loss for each sample
|
||||
for i in 0..batch_size {
|
||||
let label = labels_data[i] as usize;
|
||||
|
||||
if label >= num_labels {
|
||||
return Err(TransformerError::invalid_input(format!(
|
||||
"Label {label} exceeds number of classes {num_labels}"
|
||||
)));
|
||||
}
|
||||
|
||||
// Extract logits for this sample
|
||||
let start_idx = i * num_labels;
|
||||
let sample_logits: Vec<f32> = logits_data[start_idx..start_idx + num_labels].to_vec();
|
||||
|
||||
// Compute softmax for numerical stability
|
||||
let max_logit = sample_logits
|
||||
.iter()
|
||||
.fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_logits: Vec<f32> = sample_logits
|
||||
.iter()
|
||||
.map(|&x| (x - max_logit).exp())
|
||||
.collect();
|
||||
let sum_exp: f32 = exp_logits.iter().sum();
|
||||
|
||||
// Cross-entropy loss: -log(p_true)
|
||||
let prob = exp_logits[label] / sum_exp;
|
||||
total_loss += -prob.ln();
|
||||
}
|
||||
|
||||
// Average loss over batch
|
||||
let average_loss = total_loss / batch_size as f32;
|
||||
|
||||
Tensor::scalar(average_loss, rtx_tensor::DType::F32, logits.device())
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// BERT token classification head
|
||||
#[derive(Debug)]
|
||||
pub struct BertTokenClassificationHead {
|
||||
pub hidden_size: usize,
|
||||
pub num_labels: usize,
|
||||
pub dropout_prob: f64,
|
||||
/// Token classifier dense layer: [`hidden_size`, `num_labels`]
|
||||
pub classifier: Option<Tensor>,
|
||||
/// Classifier bias: [`num_labels`]
|
||||
pub bias: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl BertTokenClassificationHead {
|
||||
/// Create new token classification head
|
||||
pub fn new(
|
||||
hidden_size: usize,
|
||||
num_labels: usize,
|
||||
dropout_prob: f64,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
// Initialize classifier layer with Xavier initialization
|
||||
let init_std = (2.0 / (hidden_size + num_labels) as f32).sqrt();
|
||||
let classifier = Tensor::randn(&[hidden_size, num_labels], device)?.mul_scalar(init_std)?;
|
||||
let bias = Tensor::zeros([num_labels], device)?;
|
||||
|
||||
Ok(Self {
|
||||
hidden_size,
|
||||
num_labels,
|
||||
dropout_prob,
|
||||
classifier: Some(classifier),
|
||||
bias: Some(bias),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass of token classification head
|
||||
pub fn forward(&self, sequence_output: &Tensor) -> Result<Tensor> {
|
||||
let classifier_weight = self
|
||||
.classifier
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("Token classifier not initialized"))?;
|
||||
let bias = self.bias.as_ref().ok_or_else(|| {
|
||||
TransformerError::architecture("Token classifier bias not initialized")
|
||||
})?;
|
||||
|
||||
// Apply dropout (simplified - in practice would check training mode)
|
||||
// For now, skip dropout during inference
|
||||
|
||||
// Apply classifier: sequence_output @ classifier + bias
|
||||
// sequence_output shape: [batch_size, seq_len, hidden_size]
|
||||
// classifier shape: [hidden_size, num_labels]
|
||||
// output shape: [batch_size, seq_len, num_labels]
|
||||
let logits = sequence_output
|
||||
.matmul(classifier_weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
let logits = logits
|
||||
.add(bias)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
|
||||
Ok(logits)
|
||||
}
|
||||
|
||||
/// Compute token-level cross-entropy loss
|
||||
pub fn compute_loss(
|
||||
&self,
|
||||
logits: &Tensor,
|
||||
labels: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
// logits shape: [batch_size, seq_len, num_labels]
|
||||
// labels shape: [batch_size, seq_len]
|
||||
// attention_mask shape: [batch_size, seq_len] (optional)
|
||||
|
||||
let batch_size = logits.shape().dims()[0];
|
||||
let seq_len = logits.shape().dims()[1];
|
||||
let num_labels = logits.shape().dims()[2];
|
||||
|
||||
let logits_data = logits.to_cpu()?;
|
||||
let labels_data = labels.to_cpu()?;
|
||||
let mask_data = if let Some(mask) = attention_mask {
|
||||
Some(mask.to_cpu()?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut total_loss = 0.0;
|
||||
let mut valid_tokens = 0;
|
||||
|
||||
// Compute cross-entropy loss for each token position
|
||||
for batch_idx in 0..batch_size {
|
||||
for seq_idx in 0..seq_len {
|
||||
let label = labels_data[batch_idx * seq_len + seq_idx] as i32;
|
||||
|
||||
// Skip positions with -100 label (ignored tokens)
|
||||
if label == -100 {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Skip positions that are masked out
|
||||
if let Some(ref mask) = mask_data
|
||||
&& mask[batch_idx * seq_len + seq_idx] == 0.0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
valid_tokens += 1;
|
||||
|
||||
// Extract logits for this position
|
||||
let start_idx = (batch_idx * seq_len + seq_idx) * num_labels;
|
||||
let position_logits: Vec<f32> =
|
||||
logits_data[start_idx..start_idx + num_labels].to_vec();
|
||||
|
||||
// Compute softmax for numerical stability
|
||||
let max_logit = position_logits
|
||||
.iter()
|
||||
.fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_logits: Vec<f32> = position_logits
|
||||
.iter()
|
||||
.map(|&x| (x - max_logit).exp())
|
||||
.collect();
|
||||
let sum_exp: f32 = exp_logits.iter().sum();
|
||||
|
||||
// Cross-entropy loss: -log(p_true)
|
||||
if label >= 0 && (label as usize) < num_labels {
|
||||
let prob = exp_logits[label as usize] / sum_exp;
|
||||
total_loss += -prob.ln();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Average loss over valid tokens
|
||||
let average_loss = if valid_tokens > 0 {
|
||||
total_loss / valid_tokens as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
Tensor::scalar(average_loss, rtx_tensor::DType::F32, logits.device())
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// BERT question answering head for span prediction
|
||||
#[derive(Debug)]
|
||||
pub struct BertQAHead {
|
||||
pub hidden_size: usize,
|
||||
/// QA classifier: [`hidden_size`, 2] for start/end logits
|
||||
pub qa_outputs: Option<Tensor>,
|
||||
/// QA bias: [2]
|
||||
pub bias: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl BertQAHead {
|
||||
/// Create new question answering head
|
||||
pub fn new(hidden_size: usize, device: &Device) -> Result<Self> {
|
||||
// Initialize QA classifier for start/end positions
|
||||
let init_std = (2.0 / (hidden_size + 2) as f32).sqrt();
|
||||
let qa_outputs = Tensor::randn(&[hidden_size, 2], device)?.mul_scalar(init_std)?;
|
||||
let bias = Tensor::zeros([2], device)?;
|
||||
|
||||
Ok(Self {
|
||||
hidden_size,
|
||||
qa_outputs: Some(qa_outputs),
|
||||
bias: Some(bias),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass to predict start and end logits
|
||||
pub fn forward(&self, sequence_output: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let qa_weight = self
|
||||
.qa_outputs
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("QA outputs not initialized"))?;
|
||||
let bias = self
|
||||
.bias
|
||||
.as_ref()
|
||||
.ok_or_else(|| TransformerError::architecture("QA bias not initialized"))?;
|
||||
|
||||
// Apply QA classifier: sequence_output @ qa_outputs + bias
|
||||
// sequence_output shape: [batch_size, seq_len, hidden_size]
|
||||
// qa_outputs shape: [hidden_size, 2]
|
||||
// output shape: [batch_size, seq_len, 2]
|
||||
let logits = sequence_output
|
||||
.matmul(qa_weight)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
let logits = logits
|
||||
.add(bias)
|
||||
.map_err(|e| TransformerError::tensor_op(e.to_string()))?;
|
||||
|
||||
// Split into start and end logits
|
||||
let batch_size = logits.shape().dims()[0];
|
||||
let seq_len = logits.shape().dims()[1];
|
||||
|
||||
let logits_data = logits.to_cpu()?;
|
||||
let mut start_data = Vec::with_capacity(batch_size * seq_len);
|
||||
let mut end_data = Vec::with_capacity(batch_size * seq_len);
|
||||
|
||||
for batch_idx in 0..batch_size {
|
||||
for seq_idx in 0..seq_len {
|
||||
let base_idx = (batch_idx * seq_len + seq_idx) * 2;
|
||||
start_data.push(logits_data[base_idx]);
|
||||
end_data.push(logits_data[base_idx + 1]);
|
||||
}
|
||||
}
|
||||
|
||||
let start_logits = Tensor::from_data(start_data, [batch_size, seq_len], logits.device())?;
|
||||
let end_logits = Tensor::from_data(end_data, [batch_size, seq_len], logits.device())?;
|
||||
|
||||
Ok((start_logits, end_logits))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user