Still WIP

This commit is contained in:
spushk 2026-05-13 14:08:53 +05:00 • committed by spushk
commit 6d28a2f53a
Signed by: spushk
GPG key ID: 7797B761E105E8AA
4 changed files with 690 additions and 383 deletions

View file

@ -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"

View file

@ -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

View file

@ -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

View file

@ -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)