Still WIP
This commit is contained in:
parent
2109aaa16d
commit
6d28a2f53a
4 changed files with 690 additions and 383 deletions
|
|
@ -12,11 +12,11 @@ edition = "2024"
|
|||
model-ai-text.path = "./model-ai-text"
|
||||
model-ai-text.version = "=0.0.1-rc.1"
|
||||
## Wasm
|
||||
wasm-bindgen.version = "=0.2.114"
|
||||
wasm-bindgen.version = "=0.2.117"
|
||||
wasm-bindgen-futures.version = "^0.4"
|
||||
js-sys.version = "^0.3"
|
||||
web-sys.version = "^0.3"
|
||||
gloo.version = "^0.11"
|
||||
gloo.version = "^0.12"
|
||||
wasm-tracing.version = "^2"
|
||||
console_error_panic_hook.version = "^0.1"
|
||||
## Logging
|
||||
|
|
@ -33,15 +33,15 @@ tokenizers.version = "^0.22"
|
|||
tokenizers.default-features = false
|
||||
tokenizers.features = ["unstable_wasm"]
|
||||
burn.git = "https://github.com/tracel-ai/burn.git"
|
||||
burn.rev = "3a0e604cb14e445cd5a3686fbe06d0e11c76f05a"
|
||||
burn.rev = "053f030d234d0edf0c6c88935a9e7a880b1a5e0e"
|
||||
burn.default-features = false
|
||||
burn.features = ["std"]
|
||||
burn-store.git = "https://github.com/tracel-ai/burn.git"
|
||||
burn-store.rev = "3a0e604cb14e445cd5a3686fbe06d0e11c76f05a"
|
||||
burn-store.rev = "053f030d234d0edf0c6c88935a9e7a880b1a5e0e"
|
||||
burn-store.default-features = false
|
||||
burn-store.features = ["std"]
|
||||
burn-onnx.git = "https://github.com/tracel-ai/burn-onnx.git"
|
||||
burn-onnx.rev = "17a0864e02c19669b143e4227038363bb5194589"
|
||||
burn-onnx.rev = "ceed7e4ea3250dcdb9c90c38e446f30696ad2372"
|
||||
burn-onnx.default-features = false
|
||||
## HF
|
||||
hf-hub.version = "^0.5"
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ model-ai-text.optional = true
|
|||
wasm-bindgen.workspace = true
|
||||
wasm-bindgen-futures.workspace = true
|
||||
js-sys.workspace = true
|
||||
js-sys.features = ["futures"]
|
||||
gloo.workspace = true
|
||||
wasm-tracing.workspace = true
|
||||
console_error_panic_hook.workspace = true
|
||||
|
|
@ -22,7 +23,7 @@ console_error_panic_hook.workspace = true
|
|||
tracing.workspace = true
|
||||
## Inf
|
||||
burn.workspace = true
|
||||
burn.features = ["ndarray", "wgpu", "webgpu"]
|
||||
burn.features = ["ndarray", "simd", "wgpu", "webgpu"]
|
||||
burn-store.workspace = true
|
||||
burn-store.features = ["burnpack"]
|
||||
tokenizers.workspace = true
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ burn.workspace = true
|
|||
burn-store.workspace = true
|
||||
burn-store.features = ["burnpack"]
|
||||
tokenizers.workspace = true
|
||||
half.version = "^2"
|
||||
## Seri
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
use burn_onnx::ModelGen;
|
||||
use serde::{Deserialize, Deserializer};
|
||||
use std::{env, path::Path};
|
||||
use std::{
|
||||
env,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
|
|
@ -55,7 +58,7 @@ impl TokenizerConfig {
|
|||
{
|
||||
// This is configuration sanity check.
|
||||
if v > (LENGTH_CAP * 3) as u64 {
|
||||
println!(
|
||||
eprintln!(
|
||||
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
||||
TokenizerConfig::max_length_default()
|
||||
);
|
||||
|
|
@ -71,7 +74,7 @@ impl TokenizerConfig {
|
|||
{
|
||||
// This is configuration sanity check.
|
||||
if v > (LENGTH_CAP * 3) as i64 {
|
||||
println!(
|
||||
eprintln!(
|
||||
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
||||
TokenizerConfig::max_length_default()
|
||||
);
|
||||
|
|
@ -87,7 +90,7 @@ impl TokenizerConfig {
|
|||
{
|
||||
// This is configuration sanity check.
|
||||
if v > (LENGTH_CAP * 3) as f64 {
|
||||
println!(
|
||||
eprintln!(
|
||||
"WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
|
||||
TokenizerConfig::max_length_default()
|
||||
);
|
||||
|
|
@ -105,11 +108,28 @@ impl TokenizerConfig {
|
|||
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")
|
||||
.unwrap_or_else(|_| "onnx-community/e5-small-lora-ai-generated-detector-ONNX".into());
|
||||
.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")
|
||||
|
|
@ -126,9 +146,7 @@ fn main() {
|
|||
|
||||
// dwnl model
|
||||
let repo = api.model(model_repo);
|
||||
let model = repo
|
||||
.get("onnx/model.onnx")
|
||||
.expect("ai_text: model.onnx fetch");
|
||||
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")
|
||||
|
|
@ -148,6 +166,21 @@ fn main() {
|
|||
.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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue