63 lines
1.6 KiB
Rust
63 lines
1.6 KiB
Rust
|
use serde::Deserialize;
|
||
|
|
|
||
|
|
#[derive(Debug, Clone, Deserialize)]
|
||
|
|
pub struct LlamaConfig {
|
||
|
|
pub hidden_size: usize,
|
||
|
|
pub intermediate_size: usize,
|
||
|
|
pub num_hidden_layers: usize,
|
||
|
|
pub num_attention_heads: usize,
|
||
|
|
pub num_key_value_heads: usize,
|
||
|
|
pub vocab_size: usize,
|
||
|
|
pub hidden_act: String,
|
||
|
|
pub rms_norm_eps: f64,
|
||
|
|
pub rope_theta: f64,
|
||
|
|
pub max_position_embeddings: usize,
|
||
|
|
pub bos_token_id: u32,
|
||
|
|
pub eos_token_id: EosTokenId,
|
||
|
|
pub pad_token_id: u32,
|
||
|
|
pub tie_word_embeddings: bool,
|
||
|
|
pub head_dim: Option<usize>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Debug, Clone, Deserialize)]
|
||
|
|
#[serde(untagged)]
|
||
|
|
pub enum EosTokenId {
|
||
|
|
Single(u32),
|
||
|
|
Multiple(Vec<u32>),
|
||
|
|
}
|
||
|
|
|
||
|
|
impl EosTokenId {
|
||
|
|
pub fn first(&self) -> u32 {
|
||
|
|
match self {
|
||
|
|
EosTokenId::Single(id) => *id,
|
||
|
|
EosTokenId::Multiple(ids) => ids[0],
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn contains(&self, token: u32) -> bool {
|
||
|
|
match self {
|
||
|
|
EosTokenId::Single(id) => *id == token,
|
||
|
|
EosTokenId::Multiple(ids) => ids.contains(&token),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn as_slice(&self) -> &[u32] {
|
||
|
|
match self {
|
||
|
|
EosTokenId::Single(id) => std::slice::from_ref(id),
|
||
|
|
EosTokenId::Multiple(ids) => ids.as_slice(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl LlamaConfig {
|
||
|
|
pub fn head_dim(&self) -> usize {
|
||
|
|
self.head_dim.unwrap_or(self.hidden_size / self.num_attention_heads)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn from_json(path: &str) -> anyhow::Result<Self> {
|
||
|
|
let content = std::fs::read_to_string(path)?;
|
||
|
|
let config: LlamaConfig = serde_json::from_str(&content)?;
|
||
|
|
Ok(config)
|
||
|
|
}
|
||
|
|
}
|