commit f2b8bdca9856ec81d9415068a84c1d5f3bc6fb39 Author: masters Date: Wed Mar 11 02:36:36 2026 +0500 init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3a03209 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +wasm/target +wasm/Cargo.lock +wasm/inference-models/models +wasm/inference-models/pkg +extension/.output +extension/.wxt +extension/node_modules +extension/pnpm-lock.yaml +.export_venv diff --git a/Makefile.toml b/Makefile.toml new file mode 100644 index 0000000..0c3c236 --- /dev/null +++ b/Makefile.toml @@ -0,0 +1,127 @@ +[env.development] +_MK_WASMPACK_MOD = "--dev" + +[env.production] +_MK_WASMPACK_MOD = "--release" + +[tasks.release] +dependencies = [ + "build_inference_models", + "build_extension" +] + +[tasks.dev] +dependencies = [ + "build_inference_models", + "extension_npm_setup" +] +cwd = "extension" +command = "pnpm" +args = ["dev"] + +[tasks.build_extension] +dependencies = [ + "build_extension_mozf", + "zip_extension_mozf", + "build_extension_chro", + "zip_extension_chro" +] + +[tasks.zip_extension_chro] +cwd = "extension" +command = "pnpm" +args = ["zip:chrome"] + +[tasks.zip_extension_mozf] +cwd = "extension" +command = "pnpm" +args = ["zip:firefox"] + +[tasks.build_extension_chro] +dependencies = [ + "extension_npm_setup" +] +cwd = "extension" +command = "pnpm" +args = ["build:chrome"] + +[tasks.build_extension_mozf] +dependencies = [ + "extension_npm_setup" +] +cwd = "extension" +command = "pnpm" +args = ["build:firefox"] + +[tasks.extension_npm_setup] +cwd = "extension" +command = "pnpm" +args = ["install"] + +[tasks.build_inference_models] +dependencies = [ + "build_inference_models_ai_text", + "build_inference_models_ai_text_lite" +] + +[tasks.build_inference_models_ai_text] +env = {_MK_MODEL_AI_TEXT = "onnx-community/modernbert-ai-detection-raid-mage-ONNX",_MK_MODEL_AI_TEXT_ATTENTION = "yes",_MK_MODEL_AI_TEXT_TYPES = "no"} +cwd = "wasm/inference-models" +install_crate = "wasm-pack" +command = "wasm-pack" +args = [ + "build", + "${_MK_WASMPACK_MOD}", + "--reference-types", + "--target", "web", + "--out-dir", "pkg/inference-ai-text", + "--out-name", "inference-ai-text", + "--features", "ai_text" +] + +[tasks.build_inference_models_ai_text_lite] +env = {_MK_MODEL_AI_TEXT = "onnx-community/e5-small-lora-ai-generated-detector-ONNX",_MK_MODEL_AI_TEXT_ATTENTION = "yes",_MK_MODEL_AI_TEXT_TYPES = "yes"} +cwd = "wasm/inference-models" +install_crate = "wasm-pack" +command = "wasm-pack" +args = [ + "build", + "${_MK_WASMPACK_MOD}", + "--reference-types", + "--target", "web", + "--out-dir", "pkg/inference-ai-text-lite", + "--out-name", "inference-ai-text-lite", + "--features", "ai_text" +] + +[tasks.export_onnx] +dependencies = ["export_onnx_ai_text"] + +[tasks.export_onnx_ai_text] +dependencies = ["exportenv"] +command = ".export_venv/bin/optimum-cli" +args = [ + "export", "onnx", + "--model", "GeorgeDrayson/modernbert-ai-detection-raid-mage", + "--task", "text-classification", + "--dtype", "fp32", + "--opset", "18", + "--optimize", "O1", + "--monolith", + "--no-post-process", + "wasm/inference-models/models/ai_text" +] + +[tasks.exportenv] +dependencies = [ + "exportenv_create", + "exportenv_setup" +] + +[tasks.exportenv_create] +command = "python3" +args = ["-m", "venv", ".export_venv"] + +[tasks.exportenv_setup] +command = ".export_venv/bin/pip" +args = ["install", "--upgrade", "-r", "exportenv_requirements.txt"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..6fd5ae7 --- /dev/null +++ b/README.md @@ -0,0 +1,68 @@ +# firaid-moz +> I don't know why I named it this way + +A simple local webextension, which can run a bunch of small ML models to infer some properties of DOM content, such as: + +- Probability of text being AI/Human created +- Probability of image being AI/Human created + +## AI/Human classification + +> **NONE** of ML/AI/whatever models is 100% accurate by the very fundamental design. Any ML model *pretends* to do some job with certain degree of accuracy. So this extension *pretends* to know which content is AI-generated and which is not - with a certain degree of accuracy. +> +>If you don't completely trust the "AI slop" - you should not completely trust this extension also! +> +> False-positives? - Possible +> +> False-negatives? - Probably +> +> This extension - like **every other AI** - is a **tool** which aims to help you and **not** replace your own discretion. + +No, I do not hate the AI. It's cool. I hate the *slop* - when the entire "technical guide" article is nothing more then a bunch of AI ejacuation. But how do I know that slop is a slop? + +Yeah.. I have to read it. And I had _enough_ of reading. So I want to see a warning badge like "its AI" - to help me navigate through text/image content. + +AI/Human classification can be done in four different modes (sorted in order of preference for the potato hardware): + +0. _(potato)_: No classification +1. _(normal)_: Lightweight model, which can have "no decision at all" in some cases +2. _(high)_: Lightweight model still does its job, but when its decision is weak - heavy model takes the task +3. _(ultra)_: Heavy and more accurate model performs always, lightweight model is disabled + +The _normal_ mode requires ~150Mb of RAM for text classification and ~200Mb of RAM for image classification. + +The _ultra_ mode requires ~700Mb of RAM for text classification and ~500Mb of RAM for image classification. + +## `Microslop` mode + +> Slop-sloppity slop, micro slop slop slop, microslop microslop microslop + +That's easy. Turn is on and every time the page contains word "microsoft" (except for the URLs) - this extention will replace it with "microslop" + +## `Githief` mode + +You know what they do with your "licensed" and "private repository" code. + +Turn this on - and all "github" (except for urls) shall turn into "githief" + +## `Хуяндекс` mode + +At 2025 the "yandex" corporation was caught in bypassing the data permissions on Android smartphones to collect and possibly provide all kinds of your personal information to the third-parties; to track you; to deanon you; thus violating your consent and privacy. + +You can read more about it [here](https://localmess.github.io/) + +They stopped that practice after disclosure, yes, but they did it nontheless. + +So you can turn on this feature and every occurence of word "яндекс" shall be replaced for "хуяндекс" (russian wordplay, something like "fuck-yandex") and every occurence of word "yandex" (except for urls) shall be replaced with your custom word (I did not figure out proper wordplay for it yet). + +--- + +## A bit of technical details + +## Contributing + +Wooow, I am so very glad thay you're reading this section! Thank you in advance for interest to contributing. + +As you can see - this is not a github, but there is a [mirror](https://github.com/yiffyrusdev/spushks-firaid-moz), which my friend was so kind to provide. You can open issues and make PRs there - we'll merge them, preserving your identity. + +Also, if you sumbit any code, it is considered Apache-2.0 diff --git a/exportenv_requirements.txt b/exportenv_requirements.txt new file mode 100644 index 0000000..d47a54f --- /dev/null +++ b/exportenv_requirements.txt @@ -0,0 +1,2 @@ +optimum-onnx[onnxruntime] +accelerate diff --git a/extension/modules/inference-models.ts b/extension/modules/inference-models.ts new file mode 100644 index 0000000..c82c520 --- /dev/null +++ b/extension/modules/inference-models.ts @@ -0,0 +1,15 @@ +import { resolve } from "node:path"; +import { defineWxtModule } from "wxt/modules"; + +export default defineWxtModule((wxt) => { + wxt.hook("build:publicAssets", (_, assets) => { + assets.push({ + absoluteSrc: resolve("node_modules/inference-ai-text/inference-ai-text_bg.wasm"), + relativeDest: "inference-ai-text_bg.wasm", + }); + assets.push({ + absoluteSrc: resolve("node_modules/inference-ai-text-lite/inference-ai-text-lite_bg.wasm"), + relativeDest: "inference-ai-text-lite_bg.wasm", + }); + }); +}); diff --git a/extension/package.json b/extension/package.json new file mode 100644 index 0000000..c557b9c --- /dev/null +++ b/extension/package.json @@ -0,0 +1,37 @@ +{ + "name": "firaid-moz", + "version": "0.0.1", + "description": "This extension will search and highlight AI-generated text and images with small badges while you surfing. Our keen and friendly nose shall give you a heads-up, so you can dodge the slop 🦊. By the way, none data will ever leave your browser - this extension works absolutely locally. We also have cool features like 'Microslop-mode' and 'PII protection'", + "type": "module", + "scripts": { + "dev": "wxt", + "dev:firefox": "wxt -b firefox", + "build:chrome": "wxt build --mv3", + "build:firefox": "wxt build -b firefox --mv3", + "zip:chrome": "wxt zip --mv3", + "zip:firefox": "wxt zip -b firefox --mv3", + "check": "svelte-check --tsconfig ./tsconfig.json", + "postinstall": "wxt prepare" + }, + "keywords": [], + "author": "Svyatoslav Pushkarev ", + "license": "Apache-2.0", + "packageManager": "pnpm@10.28.2", + "dependencies": { + "@webext-core/messaging": "^2.3.0", + "inference-ai-text": "file:../wasm/inference-models/pkg/inference-ai-text", + "inference-ai-text-lite": "file:../wasm/inference-models/pkg/inference-ai-text-lite" + }, + "devDependencies": { + "@tailwindcss/vite": "^4.2.1", + "@tsconfig/svelte": "^5.0.7", + "@wxt-dev/module-svelte": "^2.0.4", + "daisyui": "^5.5.19", + "svelte": "^5.50.0", + "svelte-check": "^4.3.6", + "tailwindcss": "^4.2.1", + "tslib": "^2.8.1", + "typescript": "^5.9.3", + "wxt": "^0.20.14" + } +} diff --git a/extension/src/app.config.ts b/extension/src/app.config.ts new file mode 100644 index 0000000..08cfa4a --- /dev/null +++ b/extension/src/app.config.ts @@ -0,0 +1,7 @@ +import { defineAppConfig } from "#imports"; + +declare module "wxt/utils/define-app-config" { + export interface WxtAppConfig {} +} + +export default defineAppConfig({}); diff --git a/extension/src/assets/favicon.png b/extension/src/assets/favicon.png new file mode 100644 index 0000000..b89145c Binary files /dev/null and b/extension/src/assets/favicon.png differ diff --git a/extension/src/entrypoints/background.ts b/extension/src/entrypoints/background.ts new file mode 100644 index 0000000..2443c43 --- /dev/null +++ b/extension/src/entrypoints/background.ts @@ -0,0 +1,156 @@ +import inference_ai_text, { AiText, Capability as AiTextCapability } from "inference-ai-text"; +import inference_ai_text_lite, { + AiText as AiTextLite, + Capability as AiTextLiteCapability, +} from "inference-ai-text-lite"; + +import { onMessage } from "@/utils/messaging"; +import { ExtLogLevel } from "@/utils/logging"; +import { ExtensionConfigAiText } from "@/utils/config"; + +export default defineBackground({ + type: "module", + main() { + setup_inference_ai_text(); + setup_ext_console(); + }, +}); + +function setup_ext_console() { + onMessage("log", ({ data }) => { + switch (data.level) { + case ExtLogLevel.Debug: + console.debug(`[${data.sender}]`, ...data.args); + break; + case ExtLogLevel.Info: + console.log(`[${data.sender}]`, ...data.args); + break; + case ExtLogLevel.Warning: + console.warn(`[${data.sender}]`, ...data.args); + break; + case ExtLogLevel.Error: + console.error(`[${data.sender}]`, ...data.args); + break; + } + }); +} + +const AI_TEXT_STATE: { + handle: AiText | null; + queue_inference_timer: number | null; + have_normal_model: boolean; + have_lite_model: boolean; +} = { + have_normal_model: false, + have_lite_model: false, + handle: null, + queue_inference_timer: null, +}; + +function setup_inference_ai_text() { + extensionConfig.ai_text.getValue().then((config) => { + configure_ai_text(config); + }); + extensionConfig.ai_text.watch((config) => { + configure_ai_text(config); + }); + + onMessage("inference_aiText_classify", ({ data: { text } }) => { + return new Promise(async (resolve, reject) => { + if (!AI_TEXT_STATE.handle) { + reject({ + reason: "TextAi inference is disabled", + errors: [], + }); + return; + } + + console.debug("TEXT", text); + + AI_TEXT_STATE.handle + .queue(text) + .then((r) => { + // we're passing that into outer space, can not play with reference + resolve({ prob_ai: r.prob_ai, prob_human: r.prob_human }); + }) + .catch((e) => { + reject({ + reason: "inference error", + errors: [e], + }); + }); + + // Run inference if reported queue size reached threshold, + // or schedule inference to run anyway after some time + let stats = AI_TEXT_STATE.handle.stats(); + if (stats.queue_size > stats.queue_limit / 2) { + AI_TEXT_STATE.handle.run_queued_inference().catch((e) => { + console.error("Queue-threshold initiated inference failed", e); + }); + if (AI_TEXT_STATE.queue_inference_timer) { + clearTimeout(AI_TEXT_STATE.queue_inference_timer); + } + AI_TEXT_STATE.queue_inference_timer = null; + } else if (!AI_TEXT_STATE.queue_inference_timer) { + AI_TEXT_STATE.queue_inference_timer = setTimeout(() => { + if (AI_TEXT_STATE.handle) { + AI_TEXT_STATE.queue_inference_timer = null; + AI_TEXT_STATE.handle.run_queued_inference().catch((e) => { + console.error("Timed initiated inference failed", e); + }); + } + }, 500); + } + }); + }); +} + +async function configure_ai_text(config: ExtensionConfigAiText) { + if (config.enabled) { + const loader = config.lite ? inference_ai_text_lite : inference_ai_text; + const handler = config.lite ? AiTextLite : AiText; + const capability = config.lite ? AiTextLiteCapability : AiTextCapability; + const have = config.lite ? AI_TEXT_STATE.have_lite_model : AI_TEXT_STATE.have_normal_model; + const set_have = config.lite + ? () => (AI_TEXT_STATE.have_lite_model = true) + : () => (AI_TEXT_STATE.have_normal_model = true); + + if (!have) { + await loader(); + if (await guess_has_webgpu()) { + console.warn( + "WebGPU detected, but use is not currently possible due to burn APIs being sync. Check the issue progress in burn and burn-onnx repos. In the mean time will use NdArray anyway.", + ); + //await capability.assert_ndarray(); + await capability.assert_webgpu(); + } else { + await capability.assert_ndarray(); + } + set_have(); + } + + AI_TEXT_STATE.handle = new handler(); + } else { + AI_TEXT_STATE.handle = null; + if (AI_TEXT_STATE.queue_inference_timer) { + clearTimeout(AI_TEXT_STATE.queue_inference_timer); + } + } +} + +async function guess_has_webgpu() { + console.debug("Best effort to check whether WebGPU is available"); + if ("gpu" in navigator && navigator.gpu) { + console.debug("guess_webgpu: navigator has GPU"); + try { + if (await (navigator.gpu as any).requestAdapter()) { + console.debug("guess_webgpu: adapter seems fine"); + return true; + } + } catch (e) { + console.warn("guess_webgpu: adapter fail", e); + } + } + + return false; +} diff --git a/extension/src/entrypoints/content.ts b/extension/src/entrypoints/content.ts new file mode 100644 index 0000000..6dccca9 --- /dev/null +++ b/extension/src/entrypoints/content.ts @@ -0,0 +1,114 @@ +import { EXTENSION_CONFIG, ExtensionConfigAiImage, ExtensionConfigAiText } from "@/utils/config"; +import { traverse_text } from "@/utils/domtraverse"; + +export default defineContentScript({ + matches: ["http://*/*", "https://*/*"], + main() { + const trigger = setTimeout(() => { + text_ai_traverse(); + }, 1000); + + document.addEventListener("load", () => { + clearTimeout(trigger); + + text_ai_traverse(); + }); + }, +}); + +const TRAVERSE_VIEWPORT_OBSERVER = new IntersectionObserver( + (entries, observer) => { + for (const entry of entries) { + if (entry.isIntersecting) { + const stack = TRAVERSE_PARENT_TO_STACK.get(entry.target); + + if (stack) { + sendMessage("inference_aiText_classify", { text: stack.text }) + .then((result) => { + observer.unobserve(entry.target); + + if ( + EXTENSION_CONFIG.ai_text.show_ai && + result.prob_ai >= EXTENSION_CONFIG.ai_text.threshold && + result.prob_ai >= result.prob_human + ) { + ai_set_badge(entry.target, result.prob_ai, undefined); + } else if ( + EXTENSION_CONFIG.ai_text.show_human && + result.prob_human >= EXTENSION_CONFIG.ai_text.threshold && + result.prob_human >= result.prob_ai + ) { + ai_set_badge(entry.target, undefined, result.prob_human); + } else { + ai_clear_badge(entry.target); + } + }) + .catch((e) => {}); + } + } + } + }, + { + rootMargin: "500px", + threshold: 0.1, + }, +); + +interface WeakStack { + text: string; + words: number; + first_node?: Node; + last_node?: Node; + root?: Node; +} +const TRAVERSE_PARENT_TO_STACK: WeakMap = new WeakMap(); +function text_ai_traverse() { + for (const stack of traverse_text(600, 20)) { + let target = stack.first_node.parentElement || stack.root.parentElement; + if (target) { + ai_set_badge(target); + TRAVERSE_PARENT_TO_STACK.set(target, stack); + TRAVERSE_VIEWPORT_OBSERVER.observe(target); + } + } +} + +const SET_BADGES: WeakMap = new WeakMap(); +function ai_set_badge(node: Node, ai?: number, human?: number) { + const badge = document.createElement("span"); + + badge.style.display = "inline"; + badge.style.borderRadius = "4px"; + badge.style.textAlign = "center"; + badge.style.fontFamily = "sans"; + badge.style.fontSize = "12px"; + badge.style.padding = "1px"; + badge.style.marginRight = "4px"; + if (ai || human) { + if (ai) { + badge.style.backgroundColor = "#ffc36d"; + badge.style.color = "#fff"; + badge.textContent = `🤖: ${Math.trunc(ai * 100)}%`; + } else if (human) { + badge.style.backgroundColor = "#8fe1a1"; + badge.style.color = "#fff"; + badge.textContent = `👨‍💼: ${Math.trunc(human * 100)}%`; + } else { + badge.style.backgroundColor = "#ddd"; + badge.style.color = "#fff"; + badge.textContent = "🦊: confused"; + } + } else { + badge.style.backgroundColor = "#ddd"; + badge.style.color = "#fff"; + badge.textContent = "• • • • • •"; + } + + ai_clear_badge(node); + const ins_badge = node.insertBefore(badge, node.firstChild); + SET_BADGES.set(node, ins_badge); +} + +function ai_clear_badge(node: Node) { + SET_BADGES.get(node)?.remove(); +} diff --git a/extension/src/entrypoints/popup/App.svelte b/extension/src/entrypoints/popup/App.svelte new file mode 100644 index 0000000..0cb4aa2 --- /dev/null +++ b/extension/src/entrypoints/popup/App.svelte @@ -0,0 +1,223 @@ + + +
+
+
🦊
+ +
+ +
+
+ AI Text detection + + + + + + +
+
+ AI Image detection + + + + + + +
+
+ Do not check on these sites: +
+ {#each Object.entries(config_ai.value.domain_blacklist) as domain (domain[0])} +
+ {domain[0]} +
+ {/each} +
+
+
+ +
+

+ The source code of this extension is available here under terms of Apache-2.0 license +

+
+

+ The following people/organizations made this fun extention possible to create by training and + providing ML models on huggingface: +

+
+
+
+ Dugan, L., Hwang, A., Trhlı́k, F., Zhu, A., Ludan, J. M., Xu, H., Ippolito, D., & + Callison-Burch, C. (2024). RAID: A Shared Benchmark for Robust Evaluation of + Machine-Generated Text Detectors. Proceedings of the 62nd Annual Meeting of the Association for Computational + Linguistics (Volume 1: Long Papers), 12463–12492. + https://aclanthology.org/2024.acl-long.674 +
+
+
+
+
+ Drayson, G., Yilmaz, E., & Lampos, V. (2025). Machine-generated text detection + prevents language model collapse. ArXiv Preprint ArXiv:2502.15654. +
+
+
+

+ ... and many many of awesome people who develop and maintain more then 1700 cool things this + simple extension requires +

+
+

+ Let's not forget how many shoulders we're standing on 😉 +

+
+
+
diff --git a/extension/src/entrypoints/popup/app.css b/extension/src/entrypoints/popup/app.css new file mode 100644 index 0000000..4c1b0c2 --- /dev/null +++ b/extension/src/entrypoints/popup/app.css @@ -0,0 +1,2 @@ +@import "tailwindcss"; +@plugin "daisyui"; diff --git a/extension/src/entrypoints/popup/app.ts b/extension/src/entrypoints/popup/app.ts new file mode 100644 index 0000000..88fb0de --- /dev/null +++ b/extension/src/entrypoints/popup/app.ts @@ -0,0 +1,7 @@ +import { mount } from "svelte"; +import App from "./App.svelte"; +import "./app.css"; + +export default mount(App, { + target: document.getElementById("app")!, +}); diff --git a/extension/src/entrypoints/popup/index.html b/extension/src/entrypoints/popup/index.html new file mode 100644 index 0000000..986892c --- /dev/null +++ b/extension/src/entrypoints/popup/index.html @@ -0,0 +1,12 @@ + + + + + + + + +
+ + + diff --git a/extension/src/utils/config.ts b/extension/src/utils/config.ts new file mode 100644 index 0000000..e40c5e6 --- /dev/null +++ b/extension/src/utils/config.ts @@ -0,0 +1,158 @@ +import { ExtensionConfigSchema } from "#imports"; +import { defineExtensionStorage, ExtensionStorage } from "@webext-core/storage"; + +export interface ExtensionConfigAiText { + enabled: boolean; + lite: boolean; + show_human: boolean; + show_ai: boolean; + threshold: number; +} + +export interface ExtensionConfigAiImage { + enabled: boolean; + lite: boolean; + show_human: boolean; + show_ai: boolean; + threshold: number; +} + +export interface ExtensionConfigAi { + /// shall contain domain qualified name, separated by dots + /// indexed by domain name + /// e.g. { + /// "github.com": ["github", "com"] + /// } + domain_blacklist: { + [domain: string]: string[]; + }; +} + +// export interface ExtensionConfigSchema { +// version: string; +// ai_text: ExtensionConfigAiText; +// ai_image: ExtensionConfigAiImage; +// } + +//export const extensionConfig = defineExtensionStorage(browser.storage.local); + +export const extensionConfig = { + ai: storage.defineItem("local:config:ai", { + fallback: { + domain_blacklist: {}, + }, + init: () => ({ + domain_blacklist: {}, + }), + version: 1, + migrations: {}, + }), + ai_text: storage.defineItem("local:config:ai_text", { + fallback: { + enabled: true, + lite: true, + show_ai: true, + show_human: false, + threshold: 0.6, + }, + init: () => ({ + enabled: true, + lite: true, + show_ai: true, + show_human: false, + threshold: 0.6, + }), + version: 1, + migrations: {}, + }), + ai_image: storage.defineItem("local:config:ai_image", { + fallback: { + enabled: true, + lite: true, + show_ai: true, + show_human: false, + threshold: 0.6, + }, + init: () => ({ + enabled: true, + lite: true, + show_ai: true, + show_human: true, + threshold: 0.6, + }), + version: 1, + migrations: {}, + }), +} as const; + +export const EXTENSION_CONFIG: { + ai: ExtensionConfigAi; + ai_text: ExtensionConfigAiText; + ai_image: ExtensionConfigAiImage; +} = { + ai: extensionConfig.ai.fallback, + ai_text: extensionConfig.ai_text.fallback, + ai_image: extensionConfig.ai_image.fallback, +}; + +extensionConfig.ai.getValue().then((config) => { + EXTENSION_CONFIG.ai = config; + + extensionConfig.ai.watch((config) => { + EXTENSION_CONFIG.ai = config; + }); +}); + +extensionConfig.ai_text.getValue().then((config) => { + EXTENSION_CONFIG.ai_text = config; + + extensionConfig.ai_text.watch((config) => { + EXTENSION_CONFIG.ai_text = config; + }); +}); + +extensionConfig.ai_image.getValue().then((config) => { + EXTENSION_CONFIG.ai_image = config; + + extensionConfig.ai_image.watch((config) => { + EXTENSION_CONFIG.ai_image = config; + }); +}); + +// extensionConfig.getItem("version").then(async (v) => { +// const current_version = browser.runtime.getManifest().version; +// if (v != current_version) { +// switch (v) { +// default: +// await config_migrate_fresh(); +// } +// await extensionConfig.setItem("version", current_version); +// } +// }); + +// async function config_migrate_fresh() { +// await extensionConfig.clear(); +// await extensionConfig.setItem("ai_text", { +// enabled: true, +// lite: true, +// show_ai: true, +// show_human: false, +// threshold: 0.6, +// }); +// await extensionConfig.setItem("ai_image", { +// enabled: true, +// lite: true, +// show_ai: true, +// show_human: false, +// threshold: 0.6, +// }); +// await extensionConfig.setItem("fun_githief", { +// enabled: false, +// }); +// await extensionConfig.setItem("fun_microslop", { +// enabled: false, +// }); +// await extensionConfig.setItem("fun_huyandex", { +// enabled: false, +// }); +// } diff --git a/extension/src/utils/domtraverse.ts b/extension/src/utils/domtraverse.ts new file mode 100644 index 0000000..1abe88f --- /dev/null +++ b/extension/src/utils/domtraverse.ts @@ -0,0 +1,211 @@ +/// Extracts meaningful portions of text from DOM, yielding the top matched nodes +/// +/// e.g. +/// +/// ```html +///
+/// my hair are NOT +/// +/// blue! +/// +///
+/// ``` +/// +/// Shall yield a reference to the `div` with text "my hair are NOT blue!" + +const IS_DEV = import.meta.env.MODE == "development"; + +const TRAVERSE_STACK_BANS: Set = new Set([ + "IMG", + "FIGURE", + "SVG", + "VIDEO", + "META", +]); +const TRAVERSE_STACK_ROOTS: Set = new Set([ + "DIV", + "ARTICLE", + "SECTION", + "MAIN", + "HEADER", + "BLOCKQUOTE", +]); +const TRAVERSE_ALLOWED_ELEMENT_TAGS: Set = new Set([ + ...TRAVERSE_STACK_ROOTS, + "SPAN", + "P", + "H1", + "H2", + "H3", + "H4", + "H5", + "H6", + "H7", + "H8", + "UL", + "OL", + "LI", + "DEL", + "B", + "I", + "STRONG", + "EM", +]); +export function* traverse_text(soft_max_words: number, hard_min_words: number) { + const walker = document.createTreeWalker(document.body, NodeFilter.SHOW_TEXT, { + acceptNode: (node) => { + // sanity check + if (!node.isConnected) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!IS_CONNECTED"); + } + return NodeFilter.FILTER_REJECT; + } + if (node.parentNode?.nodeType != Node.ELEMENT_NODE) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!ELEMENT_NODE"); + } + return NodeFilter.FILTER_REJECT; + } + + // filter elements + if (!TRAVERSE_ALLOWED_ELEMENT_TAGS.has(node.parentElement?.tagName)) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!ALLOWED_TAG"); + } + return NodeFilter.FILTER_REJECT; + } + + // TODO: figure out more performant way to check this. + // Ignore anything inside of banned elements. + // + // For now I'll just limit the depth to 16 + let checking = node.parentElement; + for (let i = 0; i < 16; i++) { + if (!checking || checking.tagName == "BODY") { + break; + } + if (TRAVERSE_STACK_BANS.has(checking.tagName)) { + if (IS_DEV) { + node.parentElement?.setAttribute( + "data-firaid-moz-traverse", + "REJECT-BANNED_ANCESTOR", + ); + } + return NodeFilter.FILTER_REJECT; + } + checking = checking.parentElement; + } + + // drop with nonsense content + if ((node.textContent?.trim()?.length || 0) < 8) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-SHORT_TEXT"); + } + return NodeFilter.FILTER_REJECT; + } + + // finally good ones + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-traverse", "ACCEPT"); + } + return NodeFilter.FILTER_ACCEPT; + }, + }); + + let stack: { root: Node; first_node: Node; last_node: Node; text: string; words: number } | null = null; + + while (walker.nextNode()) { + const node = walker.currentNode; + + let rootable_parent = node.parentElement; + while ( + rootable_parent && + rootable_parent.tagName != "HTML" && + !TRAVERSE_STACK_ROOTS.has(rootable_parent.tagName) + ) { + rootable_parent = rootable_parent.parentElement; + } + if (!rootable_parent) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-stack", "REJECTED-!ROOTABLE_PARENT"); + } + continue; + } + + let nodeText = node.textContent?.trim(); + if (!nodeText) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-stack", "REJECTED-!TEXT"); + } + continue; + } + let nodeTextWords = count_words_naive(nodeText); + + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-stack", "VISITED"); + } + + if (!stack) { + stack = { + root: rootable_parent, + first_node: node, + last_node: node, + text: nodeText, + words: nodeTextWords, + }; + } else if (stack.root.contains(rootable_parent)) { + if (IS_DEV) { + node.parentElement?.setAttribute("data-firaid-moz-stack", "MEMBER"); + } + stack.text += " " + nodeText; + stack.words += nodeTextWords; + stack.last_node = node; + } else { + if (stack.words >= hard_min_words) { + yield stack; + if (IS_DEV) { + stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "ROOT"); + } + } else if (IS_DEV) { + stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "SKIP-SHORT_TEXT"); + } + stack = { + root: rootable_parent, + first_node: node, + last_node: node, + text: nodeText, + words: nodeTextWords, + }; + continue; + } + + if (stack.words > soft_max_words) { + if (stack.words >= hard_min_words) { + if (IS_DEV) { + stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "ROOT"); + } + } else if (IS_DEV) { + stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "SKIP-SHORT_TEXT"); + } + stack = null; + } + } +} + +function count_words_naive(text: string): number { + const is_space = /\s/; + + let count = 0; + let word_marker = false; + for (let i = 0; i < text.length; i++) { + if (is_space.test(text[i])) { + word_marker = false; + } else if (!word_marker) { + count += 1; + word_marker = true; + } + } + + return count; +} diff --git a/extension/src/utils/logging.ts b/extension/src/utils/logging.ts new file mode 100644 index 0000000..09a122f --- /dev/null +++ b/extension/src/utils/logging.ts @@ -0,0 +1,25 @@ +import { sendMessage } from "@/utils/messaging"; + +export enum ExtLogLevel { + Debug, + Info, + Warning, + Error, +} + +export const ExtConsole = (sender: string) => { + return { + debug(...args: any[]) { + sendMessage("log", { level: ExtLogLevel.Debug, sender, args }); + }, + log(...args: any[]) { + sendMessage("log", { level: ExtLogLevel.Info, sender, args }); + }, + warn(...args: any[]) { + sendMessage("log", { level: ExtLogLevel.Warning, sender, args }); + }, + error(...args: any[]) { + sendMessage("log", { level: ExtLogLevel.Error, sender, args }); + }, + }; +}; diff --git a/extension/src/utils/messaging.ts b/extension/src/utils/messaging.ts new file mode 100644 index 0000000..bf3ba85 --- /dev/null +++ b/extension/src/utils/messaging.ts @@ -0,0 +1,9 @@ +import { defineExtensionMessaging } from "@webext-core/messaging"; +import { ExtLogLevel } from "@/utils/logging"; + +interface ProtocolMap { + inference_aiText_classify(data: { text: string }): { prob_ai: number; prob_human: number }; + log(data: { level: ExtLogLevel; sender: string; args: any[] }): void; +} + +export const { sendMessage, onMessage } = defineExtensionMessaging(); diff --git a/extension/tsconfig.json b/extension/tsconfig.json new file mode 100644 index 0000000..1c23ef1 --- /dev/null +++ b/extension/tsconfig.json @@ -0,0 +1,7 @@ +{ + "extends": "./.wxt/tsconfig.json", + "compilerOptions": { + "useDefineForClassFields": true, + }, + "files": [], +} diff --git a/extension/web-ext.config.ts b/extension/web-ext.config.ts new file mode 100644 index 0000000..db94163 --- /dev/null +++ b/extension/web-ext.config.ts @@ -0,0 +1,5 @@ +import { defineWebExtConfig } from "wxt"; + +export default defineWebExtConfig({ + chromiumArgs: ["--enable-features=Vulkan"], +}); diff --git a/extension/wxt.config.ts b/extension/wxt.config.ts new file mode 100644 index 0000000..65dbef0 --- /dev/null +++ b/extension/wxt.config.ts @@ -0,0 +1,21 @@ +import { defineConfig } from "wxt"; +import tailwindcss from "@tailwindcss/vite"; + +// See https://wxt.dev/api/config.html +export default defineConfig({ + srcDir: "src", + modules: ["@wxt-dev/module-svelte"], + manifest: { + permissions: ["storage", "activeTab", "contextMenus", "pageCapture", "offscreen"], + web_accessible_resources: [ + { + matches: [""], + resources: ["/inference-ai-text_bg.wasm", "/inference-ai-text-lite_bg.wasm"], + }, + ], + homepage_url: "https://source.spushk.ru/pub/firaid-moz", + }, + vite: () => ({ + plugins: [tailwindcss()], + }), +}); diff --git a/extension/wxt.env.d.ts b/extension/wxt.env.d.ts new file mode 100644 index 0000000..1a25456 --- /dev/null +++ b/extension/wxt.env.d.ts @@ -0,0 +1 @@ +/// diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..61bbedc --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly-2026-03-01" +targets = ["wasm32-unknown-unknown"] diff --git a/wasm/Cargo.toml b/wasm/Cargo.toml new file mode 100644 index 0000000..eef176b --- /dev/null +++ b/wasm/Cargo.toml @@ -0,0 +1,60 @@ +[workspace] +resolver = "3" +rust-version = "1.94" +edition = "2024" +members = ["inference-models", "model-ai-text"] + +[workspace.dependencies] +## Workspace models +model-ai-text.path = "./model-ai-text" +model-ai-text.version = "^0.1" +## Wasm +wasm-bindgen.version = "=0.2.114" +wasm-bindgen-futures.version = "^0.4" +js-sys.version = "^0.3" +web-sys.version = "^0.3" +gloo.version = "^0.11" +wasm-tracing.version = "^2" +console_error_panic_hook.version = "^0.1" +## Logging +tracing.version = "^0.1" +tracing.features = ["max_level_debug", "release_max_level_info"] +## Seri +serde.version = "^1" +serde_json.version = "^1" +ciborium.version = "^0.2" +## Codegen +thiserror.version = "^2" +## Inf +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.default-features = false +burn-store.git = "https://github.com/tracel-ai/burn.git" +burn-store.rev = "3a0e604cb14e445cd5a3686fbe06d0e11c76f05a" +burn-store.default-features = false +burn-onnx.git = "https://github.com/tracel-ai/burn-onnx.git" +burn-onnx.rev = "17a0864e02c19669b143e4227038363bb5194589" +burn-onnx.default-features = false +## HF +hf-hub.version = "^0.5" + +[workspace.dependencies.getrandom_v4] +package = "getrandom" +version = "^0.4" +features = ["wasm_js"] + +[workspace.dependencies.getrandom_v3] +package = "getrandom" +version = '^0.3' +features = ["wasm_js"] + +[profile.release] +opt-level = 3 +lto = "fat" +codegen-units = 1 +panic = "abort" +debug = "none" +strip = "symbols" diff --git a/wasm/inference-models/.cargo/config.toml b/wasm/inference-models/.cargo/config.toml new file mode 100644 index 0000000..82e6a7f --- /dev/null +++ b/wasm/inference-models/.cargo/config.toml @@ -0,0 +1,7 @@ +[target.wasm32-unknown-unknown] +rustflags = [ + "-Ctarget-feature=+simd128,+bulk-memory,+reference-types", + '--cfg', 'getrandom_backend="wasm_js"', + "--cfg=has_std", + "-Zthreads=6" +] diff --git a/wasm/inference-models/Cargo.toml b/wasm/inference-models/Cargo.toml new file mode 100644 index 0000000..7a56efd --- /dev/null +++ b/wasm/inference-models/Cargo.toml @@ -0,0 +1,55 @@ +[package] +name = "inference-models" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +## Models +model-ai-text.workspace = true +model-ai-text.optional = true +## Wasm +wasm-bindgen.workspace = true +wasm-bindgen-futures.workspace = true +js-sys.workspace = true +gloo.workspace = true +wasm-tracing.workspace = true +console_error_panic_hook.workspace = true +## Logging +tracing.workspace = true +## Inf +burn.workspace = true +burn.features = ["ndarray", "webgpu"] +burn-store.workspace = true +burn-store.features = ["burnpack"] +tokenizers.workspace = true +## Seri +serde.workspace = true +serde_json.workspace = true +## Codegen +thiserror.workspace = true +## Wrkrnd +getrandom_v4.workspace = true +getrandom_v3.workspace = true + +[build-dependencies] +hf-hub.workspace = true +burn-onnx.workspace = true +serde_json.workspace = true + +[features] +ai_text = ["dep:model-ai-text"] + +[package.metadata.wasm-pack] +profile.release.wasm-opt = [ + "-O4", + "-ffm", + "--enable-simd", + "--enable-bulk-memory", + "--enable-reference-types", + "--enable-nontrapping-float-to-int", + "--strip-debug", + "--strip-producers", +] diff --git a/wasm/inference-models/src/capability.rs b/wasm/inference-models/src/capability.rs new file mode 100644 index 0000000..9e8acc7 --- /dev/null +++ b/wasm/inference-models/src/capability.rs @@ -0,0 +1,52 @@ +use std::sync::OnceLock; + +use burn::backend::{ndarray::NdArrayDevice, wgpu::WgpuDevice}; +use wasm_bindgen::prelude::*; + +/// This lock shall contain the prototype Device. +/// `assert_*` functions shall set it once +/// +/// If not set when check - assume default (`CubeDevice` implements Default as safe option) +pub static CUBE_DEVICE: OnceLock = OnceLock::new(); + +#[derive(Debug)] +pub enum CubeDevice { + NdArray(NdArrayDevice), + WebGpu(WgpuDevice), +} + +impl Default for CubeDevice { + fn default() -> Self { + Self::NdArray(Default::default()) + } +} + +#[wasm_bindgen(js_namespace = Capability)] +pub async fn assert_webgpu() { + if let Some(d) = CUBE_DEVICE.get() { + tracing::warn!(asserted = ?d, "CUBE_DEVICE already set. Refuse to assert webgpu."); + return; + } + + tracing::info!("Assert WebGPU-capable environment"); + let device = Default::default(); + burn::backend::wgpu::init_setup_async::( + &device, + Default::default(), + ) + .await; + + _ = CUBE_DEVICE.set(CubeDevice::WebGpu(device)).inspect_err(|e| tracing::error!(error=?e, "Burn backend was set to WGPU, but we can not testify this with CUBE_DEVICE")); +} + +#[wasm_bindgen(js_namespace = Capability)] +pub async fn assert_ndarray() { + if let Some(d) = CUBE_DEVICE.get() { + tracing::warn!(asserted = ?d, "CUBE_DEVICE already set. Refuse to assert ndarray."); + return; + } + + tracing::info!("Assert NdArray environment"); + let device = Default::default(); + _ = CUBE_DEVICE.set(CubeDevice::NdArray(device)).inspect_err(|e| tracing::error!(error=?e, "Burn backend was set to NdArray, but we can not testify this with CUBE_DEVICE")); +} diff --git a/wasm/inference-models/src/inf/ai_text.rs b/wasm/inference-models/src/inf/ai_text.rs new file mode 100644 index 0000000..3583fa8 --- /dev/null +++ b/wasm/inference-models/src/inf/ai_text.rs @@ -0,0 +1,259 @@ +use std::{ + collections::VecDeque, + sync::{Arc, LazyLock, Mutex}, +}; + +use burn::backend::{NdArray, WebGpu, ndarray::NdArrayDevice}; +use model_ai_text::*; +use wasm_bindgen::prelude::*; + +use crate::capability::{CUBE_DEVICE, CubeDevice}; + +static NDARRAY_FALLBACK_DEVICE: NdArrayDevice = NdArrayDevice::Cpu; + +#[allow(clippy::large_enum_variant)] // only used for runtime backend dispatch +enum ModelVariant { + NdArray(Model), + WebGpu(Model), +} + +impl ModelVariant { + fn new_with_cube_device() -> Result { + let device = CUBE_DEVICE.get(); + tracing::debug!("CUBE_DEVICE is {device:?} at the point of model creation"); + + match device { + Some(CubeDevice::NdArray(d)) => Model::try_load(d).map(Self::NdArray), + Some(CubeDevice::WebGpu(d)) => Model::try_load(d).map(Self::WebGpu), + None => { + tracing::warn!("No CUBE_DEVICE was asserted. Proceed with hope for NdArray"); + Model::try_load(&NDARRAY_FALLBACK_DEVICE).map(Self::NdArray) + } + } + } + + fn batch_inference(&self, texts: Vec<&str>) -> Result, ModelInferenceErr> { + match self { + Self::NdArray(m) => m.batch_inference(texts), + Self::WebGpu(m) => m.batch_inference(texts), + } + } + + fn stats(&self) -> ModelStats { + match self { + Self::NdArray(m) => m.stats(), + Self::WebGpu(m) => m.stats(), + } + } +} + +type ModelMut = Arc>>; + +static MODEL: LazyLock = LazyLock::new(|| { + tracing::debug!("Initializing AiText model with assumed CUBE_DEVICE"); + Arc::new(Mutex::new(ModelVariant::new_with_cube_device())) +}); + +static INFERENCE_QUEUE: LazyLock>>> = + LazyLock::new(|| Arc::new(Mutex::new(VecDeque::with_capacity(INFERENCE_QUEUE_SIZE)))); + +static INFERENCE_QUEUE_SIZE: usize = 128; + +static INFERENCE_BATCH_SIZE: usize = 16; + +/// The AiText is merely a handle: +/// - Shall pass the tasks into queue +/// - Shall await for task execution on demand +/// - Shall configure the model runtime (ndarray/webgpu) +/// +/// This is pretty much a wasm-bindgen struct so expect not-very-rusty function signatures. +#[wasm_bindgen] +pub struct AiText; + +#[wasm_bindgen] +impl AiText { + /// The handle shall force model load, signalizing the err + #[wasm_bindgen(constructor)] + #[allow(clippy::new_without_default)] // FIXME: introduce fail/recovery initialization + pub fn new() -> Self { + Self + } + + /// Put the new inference task into the queue. + /// Future shall be resolved when task is moved out from queue and inference is done + /// + /// To get tasks from queue and run the inference call `run_queued_inference` + pub async fn queue(&self, text: &str) -> Result { + let task_future = { + let mut queue = INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock"); + if queue.len() > INFERENCE_QUEUE_SIZE { + tracing::warn!("INFERENCE_QUEUE overflow: {INFERENCE_QUEUE_SIZE} limit reached"); + return Err(AiTextInferenceError::ModelNotLoaded); + } + + let task = AiTextInferenceTask { + text: text.into(), + async_bridge: Default::default(), + }; + let future = task.future(); + + queue.push_back(task); + + future + }; + + task_future.await + } + + pub fn stats(&self) -> AiTextStats { + let model_stats = MODEL + .lock() + .expect("MODEL lock") + .as_ref() + .inspect_err( + |e| tracing::error!(error=?e, "Model didn't load: model stats shall be defaults"), + ) + .map(|m| m.stats()) + .ok(); + + let tokenizer_max_length = model_stats + .as_ref() + .and_then(|s| s.tokenizer_max_length) + .unwrap_or_default(); + + AiTextStats { + tokenizer_max_length, + queue_limit: INFERENCE_QUEUE_SIZE, + queue_size: INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock").len(), + } + } + + /// Run inference for batch from the queue + pub async fn run_queued_inference(&self) { + let mut queue = INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock"); + + let model_lock = MODEL.lock().expect("MODEL lock"); + let Ok(model) = model_lock + .as_ref() + .inspect_err(|e| tracing::error!(error=?e, "Can not run inference: model didn't load")) + else { + return; + }; + let tasks = { + let size = queue.len().min(INFERENCE_BATCH_SIZE); + let mut tasks = Vec::with_capacity(size); + + while tasks.len() < size { + if let Some(task) = queue.pop_front() { + tasks.push(task) + } else { + break; + } + } + + tasks + }; + let texts = Vec::from_iter(tasks.iter().map(|t| t.text.as_str())); + + tracing::info!(size = texts.len(), "Request inference"); + let result = model.batch_inference(texts); + + match result { + Ok(probs) => { + tracing::info!(size = probs.len(), "Inference complete"); + probs.into_iter().zip(tasks).for_each(|(prob, task)| { + task.set(Ok(AiTextInference { + prob_ai: prob.prob_ai, + prob_human: prob.prob_human, + })); + }); + } + Err(e) => { + tracing::error!(error=?e, "Inference error"); + tasks + .into_iter() + .for_each(|t| t.set(Err(AiTextInferenceError::InferenceExecution))); + } + } + } +} + +#[wasm_bindgen] +pub struct AiTextStats { + pub tokenizer_max_length: usize, + pub queue_limit: usize, + pub queue_size: usize, +} + +#[wasm_bindgen] +pub struct AiTextInference { + pub prob_ai: f32, + pub prob_human: f32, +} + +#[wasm_bindgen] +#[derive(Debug, Clone, thiserror::Error)] +pub enum AiTextInferenceError { + #[error("Inference execution Error. Check logs for details")] + InferenceExecution, + #[error("Model is not loaded. Check logs for details")] + ModelNotLoaded, +} + +/// This shall be the promise-like thingey, holded by the queue until +/// explicitly disposed +pub struct AiTextInferenceTask { + text: String, + /// This is the field containing inference result, which does also + /// act like a future output (promise resolve) + async_bridge: Arc>, +} + +impl AiTextInferenceTask { + fn future(&self) -> InferenceTaskPromisedFuture { + InferenceTaskPromisedFuture(self.async_bridge.clone()) + } + + fn set(&self, result: Result) { + let mut bridge = self.async_bridge.lock().expect("async_bridge lock"); + bridge.result = Some(result); + + if let Some(waker) = bridge.waker.take() { + waker.wake(); + } + } +} + +#[wasm_bindgen] +#[derive(Default)] +struct InferenceTaskAsyncBridge { + result: Option>, + waker: Option, +} + +/// Future which shall resolve into promise-compatible thing, and trace the error (if any) +/// coz error gets erased into JsValue +struct InferenceTaskPromisedFuture(Arc>); + +impl Future for InferenceTaskPromisedFuture { + type Output = Result; + + fn poll( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll { + let mut inner = self.0.lock().expect("Inference future inner lock"); + + match inner.result.take() { + Some(result) => { + std::task::Poll::Ready(result.inspect_err( + |e| tracing::error!(error=?e, "Failed to perform AiText inference"), + )) + } + None => { + inner.waker = Some(cx.waker().clone()); + std::task::Poll::Pending + } + } + } +} diff --git a/wasm/inference-models/src/inf/mod.rs b/wasm/inference-models/src/inf/mod.rs new file mode 100644 index 0000000..d88e27d --- /dev/null +++ b/wasm/inference-models/src/inf/mod.rs @@ -0,0 +1,2 @@ +#[cfg(feature = "ai_text")] +pub mod ai_text; diff --git a/wasm/inference-models/src/lib.rs b/wasm/inference-models/src/lib.rs new file mode 100644 index 0000000..e0c2ec5 --- /dev/null +++ b/wasm/inference-models/src/lib.rs @@ -0,0 +1,32 @@ +#![recursion_limit = "256"] + +use wasm_bindgen::prelude::*; + +pub mod capability; +pub mod inf; + +#[wasm_bindgen] +extern "C" { + #[wasm_bindgen(js_name = log, js_namespace = console)] + pub(crate) fn jsconsole_log(s: &str); + #[wasm_bindgen(js_name = error, js_namespace = console)] + pub(crate) fn jsconsole_err(s: &str); +} + +#[wasm_bindgen(start)] +pub fn start() { + console_error_panic_hook::set_once(); + wasm_tracing::set_as_global_default(); + + let feature_report = std::format!( + "inference-models: ai-text {ai_text}", + ai_text = cfg_select! { + feature = "ai_text" => "✅", + _ => "❌" + } + ); + + tracing::info!( + "inference-models looks good;\nfeatures: {feature_report};\nThe extension source code is available at source.spushk.ru/pub/firaid-moz\nHappy browsing!" + ); +} diff --git a/wasm/model-ai-text/Cargo.toml b/wasm/model-ai-text/Cargo.toml new file mode 100644 index 0000000..3811c53 --- /dev/null +++ b/wasm/model-ai-text/Cargo.toml @@ -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 = [] diff --git a/wasm/model-ai-text/build.rs b/wasm/model-ai-text/build.rs new file mode 100644 index 0000000..9e21adf --- /dev/null +++ b/wasm/model-ai-text/build.rs @@ -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 + 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(self, v: u64) -> Result + 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(self, v: i64) -> Result + 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(self, v: f64) -> Result + 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"); +} diff --git a/wasm/model-ai-text/src/lib.rs b/wasm/model-ai-text/src/lib.rs new file mode 100644 index 0000000..d0d43b2 --- /dev/null +++ b/wasm/model-ai-text/src/lib.rs @@ -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), + #[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), +} + +#[derive(Debug, thiserror::Error)] +pub enum ModelInferenceErr { + #[error("Failed to encode inputs")] + InputEncode(Box), +} + +#[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, +} + +pub struct Model { + model: data::Model, + tokenizer: Tokenizer, + device: &'static ::Device, +} + +impl Model { + pub fn try_load(device: &'static ::Device) -> Result { + 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, 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 = + Tensor::::from_ints(ids, self.device).reshape(shape.clone()); + + let attn: Tensor = + Tensor::::from_ints(attn, self.device).reshape(shape.clone()); + + let types: Tensor = + Tensor::::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::().array_chunks::<2>().map( + |[prob_human, prob_ai]| ModelInference { + prob_ai, + prob_human, + }, + ))) + } +}