randomshit/firaid-moz
Watch
1
0
Fork
You've already forked firaid-moz
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Still WIP

Browse source
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
2 changed files with 690 additions and 383 deletions
Download patch file Download diff file Expand all files Collapse all files

1
wasm/model-ai-text/Cargo.toml
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

49
wasm/model-ai-text/build.rs
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)