Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,312 @@
|
||||
//! Transformer-based translation implementation
|
||||
|
||||
use crate::{Result, translation::AttentionConfig, translation::AttentionLayer};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Transformer configuration for translation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TransformerConfig {
|
||||
pub num_encoder_layers: usize,
|
||||
pub num_decoder_layers: usize,
|
||||
pub hidden_size: usize,
|
||||
pub num_attention_heads: usize,
|
||||
pub intermediate_size: usize,
|
||||
pub dropout_rate: f32,
|
||||
pub max_position_embeddings: usize,
|
||||
pub vocab_size: usize,
|
||||
}
|
||||
|
||||
impl Default for TransformerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
num_encoder_layers: 6,
|
||||
num_decoder_layers: 6,
|
||||
hidden_size: 512,
|
||||
num_attention_heads: 8,
|
||||
intermediate_size: 2048,
|
||||
dropout_rate: 0.1,
|
||||
max_position_embeddings: 512,
|
||||
vocab_size: 50000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Transformer encoder layer
|
||||
pub struct TransformerEncoderLayer {
|
||||
self_attention: AttentionLayer,
|
||||
feed_forward: FeedForward,
|
||||
config: TransformerConfig,
|
||||
}
|
||||
|
||||
impl TransformerEncoderLayer {
|
||||
pub fn new(config: TransformerConfig) -> Self {
|
||||
let attention_config = AttentionConfig {
|
||||
attention_type: crate::translation::attention::AttentionType::MultiHead {
|
||||
num_heads: config.num_attention_heads,
|
||||
},
|
||||
hidden_size: config.hidden_size,
|
||||
dropout_rate: config.dropout_rate,
|
||||
temperature: 1.0,
|
||||
};
|
||||
|
||||
Self {
|
||||
self_attention: AttentionLayer::new(attention_config),
|
||||
feed_forward: FeedForward::new(config.hidden_size, config.intermediate_size),
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
hidden_states: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
// Self-attention
|
||||
let (attention_output, _) = self.self_attention.forward(
|
||||
hidden_states,
|
||||
hidden_states,
|
||||
hidden_states,
|
||||
attention_mask,
|
||||
)?;
|
||||
|
||||
// Residual connection and layer norm
|
||||
let normed_attention = self.layer_norm(&(hidden_states + &attention_output)?)?;
|
||||
|
||||
// Feed-forward
|
||||
let ff_output = self.feed_forward.forward(&normed_attention)?;
|
||||
|
||||
// Residual connection and layer norm
|
||||
let output = self.layer_norm(&(&normed_attention + &ff_output)?)?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn layer_norm(&self, input: &Tensor) -> Result<Tensor> {
|
||||
// Mock layer normalization
|
||||
Ok(input.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Transformer decoder layer
|
||||
pub struct TransformerDecoderLayer {
|
||||
self_attention: AttentionLayer,
|
||||
cross_attention: AttentionLayer,
|
||||
feed_forward: FeedForward,
|
||||
config: TransformerConfig,
|
||||
}
|
||||
|
||||
impl TransformerDecoderLayer {
|
||||
pub fn new(config: TransformerConfig) -> Self {
|
||||
let attention_config = AttentionConfig {
|
||||
attention_type: crate::translation::attention::AttentionType::MultiHead {
|
||||
num_heads: config.num_attention_heads,
|
||||
},
|
||||
hidden_size: config.hidden_size,
|
||||
dropout_rate: config.dropout_rate,
|
||||
temperature: 1.0,
|
||||
};
|
||||
|
||||
Self {
|
||||
self_attention: AttentionLayer::new(attention_config.clone()),
|
||||
cross_attention: AttentionLayer::new(attention_config),
|
||||
feed_forward: FeedForward::new(config.hidden_size, config.intermediate_size),
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
hidden_states: &Tensor,
|
||||
encoder_hidden_states: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
causal_mask: Option<&Tensor>,
|
||||
) -> Result<(Tensor, Tensor)> {
|
||||
// Self-attention with causal mask
|
||||
let (self_attn_output, _) = self.self_attention.forward(
|
||||
hidden_states,
|
||||
hidden_states,
|
||||
hidden_states,
|
||||
causal_mask,
|
||||
)?;
|
||||
|
||||
let normed_self_attn = self.layer_norm(&(hidden_states + &self_attn_output)?)?;
|
||||
|
||||
// Cross-attention
|
||||
let (cross_attn_output, cross_attn_weights) = self.cross_attention.forward(
|
||||
&normed_self_attn,
|
||||
encoder_hidden_states,
|
||||
encoder_hidden_states,
|
||||
attention_mask,
|
||||
)?;
|
||||
|
||||
let normed_cross_attn = self.layer_norm(&(&normed_self_attn + &cross_attn_output)?)?;
|
||||
|
||||
// Feed-forward
|
||||
let ff_output = self.feed_forward.forward(&normed_cross_attn)?;
|
||||
let output = self.layer_norm(&(&normed_cross_attn + &ff_output)?)?;
|
||||
|
||||
Ok((output, cross_attn_weights))
|
||||
}
|
||||
|
||||
fn layer_norm(&self, input: &Tensor) -> Result<Tensor> {
|
||||
// Mock layer normalization
|
||||
Ok(input.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed-forward network
|
||||
pub struct FeedForward {
|
||||
hidden_size: usize,
|
||||
intermediate_size: usize,
|
||||
}
|
||||
|
||||
impl FeedForward {
|
||||
pub fn new(hidden_size: usize, intermediate_size: usize) -> Self {
|
||||
Self {
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
||||
// Mock feed-forward implementation
|
||||
// Would typically be: linear -> activation -> linear
|
||||
Ok(input.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Complete transformer model for translation
|
||||
pub struct TransformerTranslator {
|
||||
encoder_layers: Vec<TransformerEncoderLayer>,
|
||||
decoder_layers: Vec<TransformerDecoderLayer>,
|
||||
config: TransformerConfig,
|
||||
}
|
||||
|
||||
impl TransformerTranslator {
|
||||
pub fn new(config: TransformerConfig) -> Self {
|
||||
let encoder_layers: Vec<_> = (0..config.num_encoder_layers)
|
||||
.map(|_| TransformerEncoderLayer::new(config.clone()))
|
||||
.collect();
|
||||
|
||||
let decoder_layers: Vec<_> = (0..config.num_decoder_layers)
|
||||
.map(|_| TransformerDecoderLayer::new(config.clone()))
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
encoder_layers,
|
||||
decoder_layers,
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encode(&self, input_ids: &Tensor, attention_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let mut hidden_states = self.embed_tokens(input_ids)?;
|
||||
|
||||
for layer in &self.encoder_layers {
|
||||
hidden_states = layer.forward(&hidden_states, attention_mask)?;
|
||||
}
|
||||
|
||||
Ok(hidden_states)
|
||||
}
|
||||
|
||||
pub fn decode(
|
||||
&self,
|
||||
target_ids: &Tensor,
|
||||
encoder_hidden_states: &Tensor,
|
||||
attention_mask: Option<&Tensor>,
|
||||
) -> Result<(Tensor, Vec<Tensor>)> {
|
||||
let mut hidden_states = self.embed_tokens(target_ids)?;
|
||||
let causal_mask = self.create_causal_mask(target_ids.shape()[1])?;
|
||||
let mut cross_attention_weights = Vec::new();
|
||||
|
||||
for layer in &self.decoder_layers {
|
||||
let (new_hidden_states, cross_attn_weights) = layer.forward(
|
||||
&hidden_states,
|
||||
encoder_hidden_states,
|
||||
attention_mask,
|
||||
Some(&causal_mask),
|
||||
)?;
|
||||
hidden_states = new_hidden_states;
|
||||
cross_attention_weights.push(cross_attn_weights);
|
||||
}
|
||||
|
||||
Ok((hidden_states, cross_attention_weights))
|
||||
}
|
||||
|
||||
fn embed_tokens(&self, input_ids: &Tensor) -> Result<Tensor> {
|
||||
// Mock token embedding + positional encoding
|
||||
let batch_size = input_ids.shape()[0];
|
||||
let seq_length = input_ids.shape()[1];
|
||||
let embeddings = Tensor::randn(
|
||||
&[batch_size, seq_length, self.config.hidden_size],
|
||||
input_ids.device(),
|
||||
)?;
|
||||
Ok(embeddings)
|
||||
}
|
||||
|
||||
fn create_causal_mask(&self, seq_length: usize) -> Result<Tensor> {
|
||||
// Create lower triangular mask for causal attention
|
||||
let mut mask_data = vec![0.0f32; seq_length * seq_length];
|
||||
for i in 0..seq_length {
|
||||
for j in 0..=i {
|
||||
mask_data[i * seq_length + j] = 1.0;
|
||||
}
|
||||
}
|
||||
|
||||
let mask = Tensor::from_data(mask_data, [seq_length, seq_length], &Device::default())?;
|
||||
Ok(mask)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rtx_tensor::Device;
|
||||
|
||||
#[test]
|
||||
fn test_transformer_config() {
|
||||
let config = TransformerConfig::default();
|
||||
assert_eq!(config.num_encoder_layers, 6);
|
||||
assert_eq!(config.num_decoder_layers, 6);
|
||||
assert_eq!(config.hidden_size, 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_feed_forward() -> Result<()> {
|
||||
let ff = FeedForward::new(512, 2048);
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let input = Tensor::randn(&[1, 10, 512], &device)?;
|
||||
let output = ff.forward(&input)?;
|
||||
assert_eq!(output.shape(), input.shape());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transformer_translator() -> Result<()> {
|
||||
let config = TransformerConfig {
|
||||
hidden_size: 64,
|
||||
num_attention_heads: 4,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let translator = TransformerTranslator::new(config);
|
||||
let input_ids = Tensor::from_data(
|
||||
vec![1.0f32, 2.0, 3.0, 4.0, 5.0],
|
||||
&[1, 5],
|
||||
&Device::default(),
|
||||
)?;
|
||||
|
||||
let encoder_output = translator.encode(&input_ids, None)?;
|
||||
assert_eq!(encoder_output.shape(), &[1, 5, 64]);
|
||||
|
||||
let target_ids = Tensor::from_data(vec![1.0f32, 2.0, 3.0], &[1, 3], &Device::default())?;
|
||||
let (decoder_output, attention_weights) =
|
||||
translator.decode(&target_ids, &encoder_output, None)?;
|
||||
|
||||
assert_eq!(decoder_output.shape(), &[1, 3, 64]);
|
||||
assert!(!attention_weights.is_empty());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user