Skip to main content

entrenar/transformer/
mod.rs

1//! Transformer module with full model implementation and weight loading
2//!
3//! This module provides:
4//! - `Transformer` - Complete transformer model for language modeling
5//! - `TransformerConfig` - Model architecture configuration
6//! - `load_safetensors_weights` - Load weights from SafeTensors files
7//! - `Architecture` - Model architecture type for weight mapping
8//! - `MultiHeadAttentionWithLoRA` - Attention with deep LoRA injection
9//! - `LoRAProjection` - Linear projection with LoRA adapters
10
11mod attention;
12mod block;
13mod config;
14pub(crate) mod cuda_block;
15mod embedding;
16mod encoder;
17mod encoder_block;
18mod feedforward;
19pub mod init;
20mod model;
21mod norm;
22pub(crate) mod weights;
23#[cfg(feature = "gpu")]
24pub mod wgpu_block;
25
26pub use attention::{LoRAProjection, MultiHeadAttention, MultiHeadAttentionWithLoRA};
27pub use config::{ModelArchitecture, TransformerConfig};
28#[cfg(feature = "cuda")]
29pub use cuda_block::{
30    BlockWeights, CudaBlock, CudaGradWorkspace, CudaNf4TransformerBlock, CudaTransformerBlock,
31    GpuBlockOptimizerState,
32};
33#[cfg(not(feature = "cuda"))]
34pub use cuda_block::{CudaBlock, CudaTransformerBlock};
35#[cfg(feature = "cuda")]
36pub(crate) use cuda_block::{CudaBlockScratch, CudaLoraGradWorkspace, GpuLoraOptimizerState};
37pub use embedding::LearnedPositionEmbedding;
38pub use encoder::EncoderModel;
39pub use encoder_block::EncoderBlock;
40pub use feedforward::EncoderFeedForward;
41pub use model::Transformer;
42pub use norm::LayerNorm;
43pub use weights::{load_safetensors_weights, validate_weights, Architecture};
44#[cfg(feature = "gpu")]
45pub use wgpu_block::WgpuForwardPass;