https://www.maiachess.com/maia-worker.js

1/**
2 * Maia3 Web Worker — runs ONNX inference off the main thread.
3 *
4 * Messages FROM main thread:
5 *   { type: 'init', modelUrl, modelVersion }
6 *   { type: 'download' }
7 *   { type: 'inference', id, tokens, eloSelfs, eloOppos, batchSize }
8 *
9 * Messages TO main thread:
10 *   { type: 'status', status }
11 *   { type: 'progress', progress }
12 *   { type: 'error', message, id? }
13 *   { type: 'inference-result', id, logitsMove, logitsValue }
14 */
15
16importScripts('/ort/ort.wasm.min.js')
17
18const ORT = ort
19ORT.env.wasm.wasmPaths = '/ort/'
20
21// ── IndexedDB storage (mirrors MaiaModelStorage) ─────────────────────────────
22
23const DB_NAME = 'MaiaModels'
24const STORE_NAME = 'models'
25const MODEL_KEY = 'maia-rapid-model'
26
27function isCompatibleModelCache(data, expectedUrl, expectedVersion) {
28  return data.url === expectedUrl && data.version === expectedVersion
29}
30
31function openDB() {
32  return new Promise((resolve, reject) => {
33    const request = indexedDB.open(DB_NAME, 1)
34    request.onerror = () => reject(request.error)
35    request.onsuccess = () => resolve(request.result)
36    request.onupgradeneeded = (event) => {
37      const db = event.target.result
38      if (!db.objectStoreNames.contains(STORE_NAME)) {
39        db.createObjectStore(STORE_NAME, { keyPath: 'id' })
40      }
41    }
42  })
43}
44
45async function getCachedModel(modelUrl, modelVersion) {
46  const db = await openDB()
47  const tx = db.transaction([STORE_NAME], 'readonly')
48  const store = tx.objectStore(STORE_NAME)
49
50  const data = await new Promise((resolve, reject) => {
51    const req = store.get(MODEL_KEY)
52    req.onsuccess = () => resolve(req.result || null)
53    req.onerror = () => reject(req.error)
54  })
55
56  if (!data) return null
57
58  if (!isCompatibleModelCache(data, modelUrl, modelVersion)) {
59    const rwTx = db.transaction([STORE_NAME], 'readwrite')
60    rwTx.objectStore(STORE_NAME).delete(MODEL_KEY)
61    return null
62  }
63
64  return await data.data.arrayBuffer()
65}
66
67async function storeModel(modelUrl, modelVersion, buffer) {
68  const db = await openDB()
69  const tx = db.transaction([STORE_NAME], 'readwrite')
70  const store = tx.objectStore(STORE_NAME)
71
72  await new Promise((resolve, reject) => {
73    const req = store.put({
74      id: MODEL_KEY,
75      url: modelUrl,
76      version: modelVersion,
77      data: new Blob([buffer]),
78      timestamp: Date.now(),
79      size: buffer.byteLength,
80    })
81    req.onsuccess = () => resolve()
82    req.onerror = () => reject(req.error)
83  })
84}
85
86// ── Worker state ─────────────────────────────────────────────────────────────
87
88let session = null
89let modelUrl = null
90let modelVersion = null
91
92async function initSession(buffer) {
93  session = await ORT.InferenceSession.create(buffer)
94}
95
96// ── Message handler ──────────────────────────────────────────────────────────
97
98self.onmessage = async (e) => {
99  const msg = e.data
100
101  try {
102    switch (msg.type) {
103      case 'init': {
104        modelUrl = msg.modelUrl
105        modelVersion = msg.modelVersion
106        postMessage({ type: 'status', status: 'loading' })
107
108        const buffer = await getCachedModel(modelUrl, modelVersion)
109        if (buffer) {
110          await initSession(buffer)
111          postMessage({ type: 'status', status: 'ready' })
112        } else {
113          postMessage({ type: 'status', status: 'no-cache' })
114        }
115        break
116      }
117
118      case 'download': {
119        postMessage({ type: 'status', status: 'downloading' })
120        postMessage({ type: 'progress', progress: 0 })
121        const response = await fetch(modelUrl)
122        if (!response.ok) throw new Error('Failed to fetch model')
123
124        let buffer
125
126        if (response.body && typeof response.body.getReader === 'function') {
127          const reader = response.body.getReader()
128          const contentLength = +(response.headers.get('Content-Length') || 0)
129          const chunks = []
130          let receivedLength = 0
131          let lastReportedProgress = 0
132
133          while (true) {
134            const { done, value } = await reader.read()
135            if (done) break
136            chunks.push(value)
137            receivedLength += value.length
138
139            if (contentLength > 0) {
140              const currentProgress = Math.floor(
141                (receivedLength / contentLength) * 100,
142              )
143              if (currentProgress >= lastReportedProgress + 10) {
144                postMessage({ type: 'progress', progress: currentProgress })
145                lastReportedProgress = currentProgress
146              }
147            }
148          }
149
150          buffer = new Uint8Array(receivedLength)
151          let position = 0
152          for (const chunk of chunks) {
153            buffer.set(chunk, position)
154            position += chunk.length
155          }
156        } else {
157          buffer = new Uint8Array(await response.arrayBuffer())
158        }
159
160        await storeModel(modelUrl, modelVersion, buffer.buffer)
161        await initSession(buffer.buffer)
162        postMessage({ type: 'progress', progress: 100 })
163        postMessage({ type: 'status', status: 'ready' })
164        break
165      }
166
167      case 'inference': {
168        if (!session) {
169          postMessage({
170            type: 'error',
171            message: 'Model not initialized',
172            id: msg.id,
173          })
174          return
175        }
176
177        const { id, tokens, eloSelfs, eloOppos, batchSize } = msg
178
179        const feeds = {
180          tokens: new ORT.Tensor('float32', new Float32Array(tokens), [
181            batchSize,
182            64,
183            12,
184          ]),
185          elo_self: new ORT.Tensor('float32', new Float32Array(eloSelfs), [
186            batchSize,
187          ]),
188          elo_oppo: new ORT.Tensor('float32', new Float32Array(eloOppos), [
189            batchSize,
190          ]),
191        }
192
193        const result = await session.run(feeds)
194
195        const logitsMove = new Float32Array(result.logits_move.data)
196        const logitsValue = new Float32Array(result.logits_value.data)
197
198        postMessage(
199          {
200            type: 'inference-result',
201            id,
202            logitsMove: logitsMove.buffer,
203            logitsValue: logitsValue.buffer,
204          },
205          [logitsMove.buffer, logitsValue.buffer],
206        )
207        break
208      }
209    }
210  } catch (err) {
211    postMessage({
212      type: 'error',
213      message: err.message || 'Unknown worker error',
214      id: msg.id,
215    })
216  }
217}