use super::TransformerConfig;
pub fn tiny_transformer(vocab_size: usize, context_len: usize) -> TransformerConfig {
TransformerConfig {
vocab_size,
context_len,
hidden_size: 128,
ffw_size: 256,
num_layers: 2,
num_heads: 4,
}
}
pub fn small_transformer(vocab_size: usize, context_len: usize) -> TransformerConfig {
TransformerConfig {
vocab_size,
context_len,
hidden_size: 256,
ffw_size: 1024,
num_layers: 6,
num_heads: 8,
}
}
pub fn base_transformer(vocab_size: usize, context_len: usize) -> TransformerConfig {
TransformerConfig {
vocab_size,
context_len,
hidden_size: 768,
ffw_size: 3072,
num_layers: 12,
num_heads: 12,
}
}