use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
pub input_names: Vec<String>,
pub output_name: String,
pub max_sequence_length: usize,
pub vocab_size: usize,
pub model_type: String,
}
impl Default for Config {
fn default() -> Self {
Self {
input_names: vec!["input_ids".to_string()],
output_name: "logits".to_string(),
max_sequence_length: 128,
vocab_size: 32000,
model_type: "coreml".to_string(),
}
}
}
impl Config {
pub fn bert_config(output_name: &str, max_seq_len: usize, vocab_size: usize) -> Self {
Self {
input_names: vec![
"input_ids".to_string(),
"token_type_ids".to_string(),
"attention_mask".to_string(),
],
output_name: output_name.to_string(),
max_sequence_length: max_seq_len,
vocab_size,
model_type: "bert".to_string(),
}
}
}