Spaces:
Running
Running
| import { env, PreTrainedModel, AutoTokenizer, Tensor } from '@huggingface/transformers'; | |
| export const MODEL_ID = 'raincandy-u/MacroStories'; | |
| export const REVISION = 'dc1ef998bf5331f7f70a77c815ed3e4f643a0ba7'; | |
| export const MODEL_PATH = 'macrostories'; | |
| env.allowRemoteModels = false; | |
| env.allowLocalModels = true; | |
| // Keep this an origin-relative path: Transformers.js metadata discovery treats | |
| // absolute HTTP paths as remote even when they refer to our own static assets. | |
| env.localModelPath = new URL('../models/', import.meta.url).pathname; | |
| env.backends.onnx.wasm.wasmPaths = new URL('../vendor/', import.meta.url).href; | |
| // A static HF Space iframe doesn't need SharedArrayBuffer or COOP/COEP headers. | |
| env.backends.onnx.wasm.numThreads = 1; | |
| env.backends.onnx.wasm.proxy = false; | |
| export async function loadRuntime(device = 'wasm', progress = () => {}) { | |
| if (!['wasm', 'webgpu'].includes(device)) throw new Error('Choose WebAssembly or WebGPU.'); | |
| if (device === 'webgpu') { | |
| if (!globalThis.navigator?.gpu || !await navigator.gpu.requestAdapter()) { | |
| throw new Error('No WebGPU adapter available.'); | |
| } | |
| } | |
| const [tokenizer, vocabulary, provenance] = await Promise.all([ | |
| AutoTokenizer.from_pretrained(MODEL_PATH, { local_files_only: true }), | |
| fetch(new URL('../models/macrostories/tokenizer.json', import.meta.url)).then(r => { if (!r.ok) throw new Error('Tokenizer file missing.'); return r.json(); }), | |
| fetch(new URL('../models/macrostories/provenance.json', import.meta.url)).then(r => { if (!r.ok) throw new Error('Model provenance missing.'); return r.json(); }), | |
| ]); | |
| // Fallback belongs to BackendController: never reuse ORT after GPU failure. | |
| const runtime = await createRuntime(device, progress); | |
| return { ...runtime, tokenizer, vocabulary: vocabulary.model.vocab, provenance }; | |
| } | |
| async function createRuntime(device, progress) { | |
| progress({ status: 'runtime', device }); | |
| // Ouro isn't a registered AutoModelForCausalLM architecture. The public base | |
| // class loads the exported graph; our loop supplies its custom four-pass KV cache. | |
| const model = await PreTrainedModel.from_pretrained(MODEL_PATH, { | |
| device, dtype: 'fp32', local_files_only: true, progress_callback: progress, | |
| // Keep the growing KV cache on the GPU. Only the 378 logits need readback | |
| // for sampling; copying all past keys/values out and back per token is wasteful. | |
| session_options: device === 'webgpu' ? { preferredOutputLocation: { present_key_values: 'gpu-buffer', logits: 'cpu' } } : {}, | |
| }); | |
| const decoder = new CachedDecoder(model); | |
| try { | |
| const logits = await decoder.forward(1); | |
| if (logits.length !== 378 || !logits.every(Number.isFinite)) throw new Error('Runtime produced invalid logits.'); | |
| decoder.reset(); | |
| } catch (error) { | |
| await model.dispose(); | |
| throw error; | |
| } | |
| return { model, decoder, device }; | |
| } | |
| export class CachedDecoder { | |
| constructor(model) { this.model = model; this.reset(); } | |
| reset() { | |
| this.past?.dispose(); | |
| this.position = 0; | |
| this.past = new Tensor('float32', new Float32Array(0), [4, 2, 0, 16]); | |
| } | |
| async forward(token) { | |
| if (!Number.isInteger(token) || token < 0 || token >= 378) throw new Error('Invalid token ID.'); | |
| if (this.position >= 2048) throw new Error('The 2,048-token context is full.'); | |
| const input = new Tensor('int64', BigInt64Array.of(BigInt(token)), [1]); | |
| const position = new Tensor('int64', BigInt64Array.of(BigInt(this.position)), [1]); | |
| let output; | |
| try { | |
| output = await this.model({ input_ids: input, position, past_key_values: this.past }); | |
| const logits = Float32Array.from(output.logits.data); | |
| this.past.dispose(); | |
| this.past = output.present_key_values; | |
| this.position++; | |
| return logits; | |
| } finally { | |
| input.dispose(); | |
| position.dispose(); | |
| output?.logits.dispose(); | |
| } | |
| } | |
| } | |
| export function encodePrompt(tokenizer, text) { | |
| const canonical = text.replace(/\b(Alex|Robin|Casey)\b/g, name => `[PERSON_NAME_${['Alex', 'Robin', 'Casey'].indexOf(name) + 1}]`); | |
| const ids = Array.from(tokenizer(canonical).input_ids.data, Number); | |
| return { ids, unknown: ids.filter(id => id === 3).length, canonical }; | |
| } | |
| export function renderStory(tokenizer, ids) { | |
| const display = ids.filter(id => ![0, 1, 2].includes(id)); | |
| if (!display.length) return { raw: '', text: '' }; | |
| const raw = tokenizer.decode(display, { skip_special_tokens: false }).trim(); | |
| const text = raw.replace(/\[PERSON_NAME_(\d+)\]/g, (match, n) => ['Alex', 'Robin', 'Casey'][Number(n) - 1] ?? match); | |
| return { raw, text }; | |
| } | |