Arx 0.5.0: Fix problems on Mac (fix.sh + fixes.json, undo), v0.6 for 4-bit builds, website wording
b2c9ced verified | import { ANSWER_SCHEMAS } from './auditor-core.js'; | |
| // Loads the auditor model in WebLLM (runs on the device's GPU through WebGPU), for the desktop app and the phone page. | |
| // The model is downloaded once and cached by the browser, so later runs work offline. | |
| // The fine-tuned Devseis auditor model (https://huggingface.co/Devseis/endpoint-auditor-0.5b), published per version | |
| // as WebLLM/MLC weights. It has the same architecture as Qwen2.5-0.5B-Instruct, so WebLLM's prebuilt runtime library | |
| // for that model is reused and only the weights differ. Set MODEL_VERSION to null to use the base model. | |
| export const MODEL_VERSION = 'v0.3'; | |
| const MODEL_REPO = (quant, version) => `https://huggingface.co/Devseis/endpoint-auditor-0.5b-${quant}-MLC/resolve/${version}/`; | |
| // Builds: q0f16 = unquantised (about 1 GB, best quality; desktop), q4f16_1 = 4-bit (about 280 MB; phones), | |
| // q4f32_1 = 4-bit for GPUs without 16-bit maths (shader-f16). | |
| const BASE_MODEL = quant => `Qwen2.5-0.5B-Instruct-${quant}-MLC`; | |
| // Output length per task (the AI-tool classification is the longest answer). | |
| const MAX_TOKENS = { finding: 900, summary: 1200, ai_classification: 1600 }; | |
| // webllm: the @mlc-ai/web-llm module (bundled in the app, from a CDN on the phone page). | |
| // precision: 'full' (desktop: q0f16 when the GPU supports 16-bit maths) or 'compact' (phones: 4-bit). | |
| // version: model tag to load (from app/model-version.json on the Space), or { version, version_4bit } to use a different | |
| // tag for the 4-bit builds; defaults to MODEL_VERSION. | |
| export async function loadModelWith(webllm, onProgress, precision = 'compact', version = MODEL_VERSION) { | |
| const { CreateMLCEngine, prebuiltAppConfig } = webllm; | |
| if (!('gpu' in navigator)) throw new Error('WebGPU is not available on this device.'); | |
| const adapter = await navigator.gpu.requestAdapter(); | |
| if (!adapter) throw new Error('No WebGPU-capable graphics adapter was found.'); | |
| const f16 = adapter.features.has('shader-f16'); | |
| const quant = !f16 ? 'q4f32_1' : precision === 'full' ? 'q0f16' : 'q4f16_1'; | |
| if (version && typeof version === 'object') version = (quant.startsWith('q4') ? version.version_4bit : null) || version.version || MODEL_VERSION; | |
| const base = prebuiltAppConfig.model_list.find(m => m.model_id === BASE_MODEL(quant)); | |
| if (!base) throw new Error('The base model is not in this WebLLM version.'); | |
| const custom = version ? { ...base, model: MODEL_REPO(quant, version), model_id: `devseis-auditor-0.5b-${version}-${quant}` } : null; | |
| const appConfig = { ...prebuiltAppConfig, model_list: custom ? [...prebuiltAppConfig.model_list, custom] : prebuiltAppConfig.model_list }; | |
| const initProgressCallback = report => onProgress({ progress: report.progress, text: report.text }); | |
| let engine, modelId, fineTuned = false; | |
| if (custom) { | |
| try { | |
| engine = await CreateMLCEngine(custom.model_id, { appConfig, initProgressCallback }); | |
| modelId = custom.model_id; | |
| fineTuned = true; | |
| } catch (error) { | |
| console.warn(`Devseis model ${version} could not be loaded (${error.message}); using the base model.`); | |
| } | |
| } | |
| if (!engine) { | |
| engine = await CreateMLCEngine(base.model_id, { appConfig, initProgressCallback }); | |
| modelId = base.model_id; | |
| } | |
| return { | |
| name: modelId, | |
| fineTuned, | |
| // Constrained decoding: the reply is always valid JSON in the answer schema of the task. | |
| async generate(system, prompt, task) { | |
| const reply = await engine.chat.completions.create({ | |
| messages: [{ role: 'system', content: system }, { role: 'user', content: prompt }], | |
| temperature: 0, | |
| max_tokens: MAX_TOKENS[task] || 900, | |
| response_format: { type: 'json_object', schema: JSON.stringify(ANSWER_SCHEMAS[task]) }, | |
| }); | |
| return reply.choices[0].message.content; | |
| }, | |
| }; | |
| } | |