236 lines
8.6 KiB
Rust
236 lines
8.6 KiB
Rust
use burn_onnx::ModelGen;
|
|
use serde::{Deserialize, Deserializer};
|
|
use std::{
|
|
env,
|
|
path::{Path, PathBuf},
|
|
};
|
|
use tokenizers::Tokenizer;
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct Config {
|
|
hidden_size: usize,
|
|
num_attention_heads: usize,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct TokenizerConfig {
|
|
pad_token: String,
|
|
#[serde(
|
|
alias = "model_max_length",
|
|
default = "TokenizerConfig::max_length_default",
|
|
deserialize_with = "TokenizerConfig::max_length_de"
|
|
)]
|
|
max_length: usize,
|
|
}
|
|
|
|
impl TokenizerConfig {
|
|
const fn max_length_default() -> usize {
|
|
512
|
|
}
|
|
|
|
/// If value is far above `LENGTH_CAP` - that means broken config. We're going to
|
|
/// play safe and replace it with `max_length_default`.
|
|
///
|
|
/// Yeah yeah, I know, "(in)sane defaults".
|
|
///
|
|
/// Yet our current model BTW (!) has insanely large `model_max_length` while
|
|
/// in reality its cap is 512.
|
|
fn max_length_de<'de, D>(deserializer: D) -> Result<usize, D::Error>
|
|
where
|
|
D: Deserializer<'de>,
|
|
{
|
|
struct Visitor;
|
|
const LENGTH_CAP: usize = 8192;
|
|
|
|
impl<'de> serde::de::Visitor<'de> for Visitor {
|
|
type Value = usize;
|
|
|
|
fn expecting(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
|
|
write!(
|
|
fmt,
|
|
"unsigned integer. Values are capped by {LENGTH_CAP}. Float is accepted, and also capped to {LENGTH_CAP}."
|
|
)
|
|
}
|
|
|
|
fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
|
|
where
|
|
E: serde::de::Error,
|
|
{
|
|
// This is configuration sanity check.
|
|
if v > (LENGTH_CAP * 3) as u64 {
|
|
eprintln!(
|
|
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
|
TokenizerConfig::max_length_default()
|
|
);
|
|
Ok(TokenizerConfig::max_length_default())
|
|
} else {
|
|
Ok(v.clamp(0, LENGTH_CAP as u64) as usize)
|
|
}
|
|
}
|
|
|
|
fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
|
|
where
|
|
E: serde::de::Error,
|
|
{
|
|
// This is configuration sanity check.
|
|
if v > (LENGTH_CAP * 3) as i64 {
|
|
eprintln!(
|
|
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
|
TokenizerConfig::max_length_default()
|
|
);
|
|
Ok(TokenizerConfig::max_length_default())
|
|
} else {
|
|
Ok(v.clamp(0, LENGTH_CAP as i64) as usize)
|
|
}
|
|
}
|
|
|
|
fn visit_f64<E>(self, v: f64) -> Result<Self::Value, E>
|
|
where
|
|
E: serde::de::Error,
|
|
{
|
|
// This is configuration sanity check.
|
|
if v > (LENGTH_CAP * 3) as f64 {
|
|
eprintln!(
|
|
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
|
TokenizerConfig::max_length_default()
|
|
);
|
|
Ok(TokenizerConfig::max_length_default())
|
|
} else {
|
|
Ok(v.clamp(0.0, LENGTH_CAP as f64) as usize)
|
|
}
|
|
}
|
|
}
|
|
|
|
deserializer.deserialize_any(Visitor)
|
|
}
|
|
}
|
|
|
|
fn main() {
|
|
println!("cargo:rerun-if-changed=build.rs");
|
|
println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT");
|
|
println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_ONNX");
|
|
println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_ATTENTION");
|
|
println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_TYPES");
|
|
|
|
let model_repo = std::env::var("_MK_MODEL_AI_TEXT")
|
|
.inspect_err(|_| {
|
|
eprintln!("_MK_MODEL_AI_TEXT env variable required to point to HF model repository")
|
|
})
|
|
.expect("_MK_MODEL_AI_TEXT is set to huggingface model repo");
|
|
let model_file =
|
|
std::env::var("_MK_MODEL_AI_TEXT_ONNX").inspect_err(|_| {
|
|
eprintln!("_MK_MODEL_AI_TEXT_ONNX env variable required to point to onnx file under HF model repository");
|
|
}).expect("_MK_MODEL_AI_TEXT_FILE is set to .onnx path at huggingface model repo");
|
|
|
|
let model_name = model_file
|
|
.rsplit_once('/')
|
|
.map(|(_, m)| m)
|
|
.unwrap_or(model_file.as_str())
|
|
.rsplit_once('.')
|
|
.map(|(name, _)| name)
|
|
.unwrap_or(model_file.as_str());
|
|
|
|
let has_attention = std::env::var("_MK_MODEL_AI_TEXT_ATTENTION")
|
|
.is_ok_and(|v| v == "true" || v == "yes" || v == "1");
|
|
let has_types = std::env::var("_MK_MODEL_AI_TEXT_TYPES")
|
|
.is_ok_and(|v| v == "true" || v == "yes" || v == "1");
|
|
|
|
if has_attention {
|
|
println!("cargo:rustc-cfg=feature=\"_has_attention\"");
|
|
}
|
|
if has_types {
|
|
println!("cargo:rustc-cfg=feature=\"_has_types\"");
|
|
}
|
|
|
|
let api = hf_hub::api::sync::Api::new().expect("Huggingface API");
|
|
|
|
// dwnl model
|
|
let repo = api.model(model_repo);
|
|
let model = repo.get(&model_file).expect("ai_text: model.onnx fetch");
|
|
let config_json = repo.get("config.json").expect("ai_text: config.json fetch");
|
|
let tokenizer_json = repo
|
|
.get("tokenizer.json")
|
|
.expect("ai_text: tokenizer.json fetch");
|
|
let tokenizer_config_json = repo
|
|
.get("tokenizer_config.json")
|
|
.expect("ai_text: tokenizer_config.json fetch");
|
|
|
|
let burned_dir = Path::new("burned").join("ai_text");
|
|
let out_dir = Path::new(&env::var("OUT_DIR").expect("OUT_DIR is env var")).join(&burned_dir);
|
|
|
|
// handle model embedding
|
|
ModelGen::new()
|
|
.load_strategy(burn_onnx::LoadStrategy::Embedded)
|
|
.simplify(true)
|
|
.input(model.to_str().unwrap())
|
|
.out_dir(burned_dir.to_str().unwrap())
|
|
.run_from_script();
|
|
|
|
// They put model into <filename>.rs, which depends on env variables set, to we move
|
|
// model into well-known location.
|
|
let loc = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR"))
|
|
.join(&burned_dir)
|
|
.join(model_name)
|
|
.with_extension("rs");
|
|
let to = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR"))
|
|
.join(&burned_dir)
|
|
.join("model.rs");
|
|
|
|
std::fs::copy(&loc, &to)
|
|
.and_then(|_| std::fs::remove_file(&loc))
|
|
.inspect_err(|e| eprintln!("While moving {loc:?} to {to:?}: {e:?}"))
|
|
.expect("To move <model_name>.rs to model.rs");
|
|
|
|
// construct tokenizer AoT
|
|
let tokenizer = {
|
|
let mut tokenizer = Tokenizer::from_file(tokenizer_json)
|
|
.inspect_err(|e| eprintln!("While loading tokenizer.json: {e:?}"))
|
|
.expect("Tokenizer JSON at tokenizer.json");
|
|
|
|
let tokenizer_config: TokenizerConfig = serde_json::from_reader(
|
|
std::fs::File::open(tokenizer_config_json)
|
|
.inspect_err(|e| eprintln!("While opening tokenizer_config.json: {e:?}"))
|
|
.expect("File tokenizer_config.json"),
|
|
)
|
|
.inspect_err(|e| eprint!("While deserializing tokenizer_config.json: {e:?}"))
|
|
.expect("Valid JSON at tokenizer_config.json");
|
|
|
|
let pad_id = tokenizer
|
|
.token_to_id(&tokenizer_config.pad_token)
|
|
.expect("Tokenizer to know PAD token");
|
|
|
|
let config: Config = serde_json::from_reader(
|
|
std::fs::File::open(config_json)
|
|
.inspect_err(|e| eprintln!("While opening config.json: {e:?}"))
|
|
.expect("Open file config.json"),
|
|
)
|
|
.inspect_err(|e| eprintln!("While deserializing config.json: {e:?}"))
|
|
.expect("Valid JSON at config.json");
|
|
|
|
tokenizer
|
|
.with_truncation(Some(tokenizers::TruncationParams {
|
|
direction: tokenizers::TruncationDirection::Right,
|
|
max_length: tokenizer_config.max_length,
|
|
strategy: Default::default(),
|
|
stride: Default::default(),
|
|
}))
|
|
.inspect_err(|e| eprintln!("While setting tokenizer truncation: {e:?}"))
|
|
.expect("Set tokenizer truncation");
|
|
|
|
tokenizer.with_padding(Some(tokenizers::PaddingParams {
|
|
strategy: tokenizers::PaddingStrategy::BatchLongest,
|
|
direction: tokenizers::PaddingDirection::Right,
|
|
pad_to_multiple_of: Some(config.hidden_size / config.num_attention_heads),
|
|
pad_id,
|
|
pad_type_id: 0,
|
|
pad_token: tokenizer_config.pad_token,
|
|
}));
|
|
|
|
tokenizer
|
|
};
|
|
|
|
tokenizer
|
|
.save(out_dir.join("tokenizer_d.json"), false)
|
|
.inspect_err(|e| eprintln!("While saving tokenizer: {e:?}"))
|
|
.expect("Save tokenizer");
|
|
}
|