init
This commit is contained in:
masters
commit
f2b8bdca98
GPG key ID:
7797B761E105E8AA
3 changed files with 2068 additions and 0 deletions
31
wasm/model-ai-text/Cargo.toml
Normal file
31
wasm/model-ai-text/Cargo.toml
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
[package]
|
||||
name = "model-ai-text"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
## Logging
|
||||
tracing.workspace = true
|
||||
## Inf
|
||||
burn.workspace = true
|
||||
burn-store.workspace = true
|
||||
burn-store.features = ["burnpack"]
|
||||
tokenizers.workspace = true
|
||||
## Seri
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
ciborium.workspace = true
|
||||
## Codegen
|
||||
thiserror.workspace = true
|
||||
|
||||
[build-dependencies]
|
||||
hf-hub.workspace = true
|
||||
burn-onnx.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokenizers.workspace = true
|
||||
ciborium.workspace = true
|
||||
|
||||
[features]
|
||||
_has_attention = []
|
||||
_has_types = []
|
||||
203
wasm/model-ai-text/build.rs
Normal file
203
wasm/model-ai-text/build.rs
Normal file
|
|
@ -0,0 +1,203 @@
|
|||
use burn_onnx::ModelGen;
|
||||
use serde::{Deserialize, Deserializer};
|
||||
use std::{env, path::Path};
|
||||
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 {
|
||||
println!(
|
||||
"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 {
|
||||
println!(
|
||||
"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 {
|
||||
println!(
|
||||
"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_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());
|
||||
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("onnx/model.onnx")
|
||||
.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();
|
||||
|
||||
// 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");
|
||||
}
|
||||
148
wasm/model-ai-text/src/lib.rs
Normal file
148
wasm/model-ai-text/src/lib.rs
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
#![recursion_limit = "256"]
|
||||
#![feature(iter_array_chunks)]
|
||||
|
||||
use burn::{prelude::*, tensor::activation::quiet_softmax};
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
mod data {
|
||||
#![allow(clippy::all)] // have massive codegen here.
|
||||
extern crate alloc;
|
||||
|
||||
include!(concat!(env!("OUT_DIR"), "/burned/ai_text/model.rs"));
|
||||
|
||||
pub static TOKENIZER: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/burned/ai_text/tokenizer_d.json"));
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ModelErr {
|
||||
#[error("Failed to load TOKENIZER")]
|
||||
Tokenizer(Box<dyn core::error::Error + Send + Sync>),
|
||||
#[error("Invalid TOKENIZER_CONFIG")]
|
||||
TokenizerConfig(serde_json::Error),
|
||||
#[error("Invalid CONFIG")]
|
||||
Config(serde_json::Error),
|
||||
#[error("Missing required token `{0}` (this is name, not the token itself)")]
|
||||
MissingToken(&'static str),
|
||||
#[error("Failed to apply truncation to the tokenizer")]
|
||||
TokenizerWithTruncation(Box<dyn core::error::Error + Send + Sync>),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ModelInferenceErr {
|
||||
#[error("Failed to encode inputs")]
|
||||
InputEncode(Box<dyn core::error::Error + Send + Sync>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ModelInference {
|
||||
pub prob_ai: f32,
|
||||
pub prob_human: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ModelStats {
|
||||
pub tokenizer_max_length: Option<usize>,
|
||||
}
|
||||
|
||||
pub struct Model<B: Backend> {
|
||||
model: data::Model<B>,
|
||||
tokenizer: Tokenizer,
|
||||
device: &'static <B as Backend>::Device,
|
||||
}
|
||||
|
||||
impl<B: Backend> Model<B> {
|
||||
pub fn try_load(device: &'static <B as Backend>::Device) -> Result<Self, ModelErr> {
|
||||
tracing::debug!(device = ?device, "Requested to load model and tokenizer");
|
||||
|
||||
let tokenizer = Tokenizer::from_bytes(data::TOKENIZER).map_err(ModelErr::Tokenizer)?;
|
||||
let model = data::Model::from_embedded(device);
|
||||
|
||||
Ok(Self {
|
||||
device,
|
||||
model,
|
||||
tokenizer,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn stats(&self) -> ModelStats {
|
||||
let tokenizer_max_length = self.tokenizer.get_truncation().map(|t| t.max_length);
|
||||
|
||||
ModelStats {
|
||||
tokenizer_max_length,
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub fn batch_inference(
|
||||
&self,
|
||||
texts: Vec<&str>,
|
||||
) -> Result<Vec<ModelInference>, ModelInferenceErr> {
|
||||
tracing::debug!(len = texts.len(), "Requested batch inference");
|
||||
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let encodings = self
|
||||
.tokenizer
|
||||
.encode_batch(texts, true)
|
||||
.map_err(ModelInferenceErr::InputEncode)?;
|
||||
|
||||
let batch_len = encodings.len();
|
||||
let encoding_len = encodings.first().map(|f| f.len()).unwrap_or_default();
|
||||
|
||||
let (ids, attn, types) = {
|
||||
let bulk_size = batch_len * encoding_len;
|
||||
|
||||
let mut buf = vec![0u32; bulk_size * 3]; // 3 bulks: ids, attn, types
|
||||
|
||||
let ids_i = 0;
|
||||
let attn_i = ids_i + bulk_size;
|
||||
let types_i = attn_i + bulk_size;
|
||||
|
||||
for (i, e) in encodings.iter().enumerate() {
|
||||
let off = i * encoding_len;
|
||||
buf[ids_i + off..ids_i + off + encoding_len].copy_from_slice(e.get_ids());
|
||||
buf[attn_i + off..attn_i + off + encoding_len]
|
||||
.copy_from_slice(e.get_attention_mask());
|
||||
buf[types_i + off..types_i + off + encoding_len].copy_from_slice(e.get_type_ids());
|
||||
}
|
||||
|
||||
let shape = Shape::new([batch_len, encoding_len]);
|
||||
|
||||
let ids = &buf[ids_i..ids_i + bulk_size];
|
||||
let attn = &buf[attn_i..attn_i + bulk_size];
|
||||
let types = &buf[types_i..types_i + bulk_size];
|
||||
|
||||
let ids: Tensor<B, 2, Int> =
|
||||
Tensor::<B, 1, Int>::from_ints(ids, self.device).reshape(shape.clone());
|
||||
|
||||
let attn: Tensor<B, 2, Int> =
|
||||
Tensor::<B, 1, Int>::from_ints(attn, self.device).reshape(shape.clone());
|
||||
|
||||
let types: Tensor<B, 2, Int> =
|
||||
Tensor::<B, 1, Int>::from_ints(types, self.device).reshape(shape);
|
||||
|
||||
(ids, attn, types)
|
||||
};
|
||||
|
||||
tracing::debug!("Batch inference ready, encoded");
|
||||
let logits = cfg_select! {
|
||||
all(feature = "_has_attention", feature = "_has_types") => self.model.forward(ids, attn, types),
|
||||
feature = "_has_types" => self.model.forward(ids, types),
|
||||
feature = "_has_attention" => self.model.forward(ids, attn),
|
||||
_ => self.model.forward(ids)
|
||||
};
|
||||
|
||||
tracing::debug!(logits = ?logits, "Batch inference complete");
|
||||
let probs = quiet_softmax(logits, 1).into_data();
|
||||
|
||||
Ok(Vec::from_iter(probs.iter::<f32>().array_chunks::<2>().map(
|
||||
|[prob_human, prob_ai]| ModelInference {
|
||||
prob_ai,
|
||||
prob_human,
|
||||
},
|
||||
)))
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue