chore: developing the package

Author: jortsupetterson

embedder/index.js

@@ -1,17 +1,14 @@
 import { mkdir, readFile, writeFile } from 'node:fs/promises'
 import { dirname } from 'node:path'
+import { toCompressed, toBase64UrlString } from '@z-base/bytecodec'
 
 const js = String.raw
 
 const modelPath = './models/quantized_models/eng/model.int4.onnx'
 const modelDataPath = './models/quantized_models/eng/model.int4.onnx.data'
 const tokenizerModelPath = './models/quantized_models/eng/tokenizer.model'
 
-const outPath = './src/Model/class.ts'
-
-function toUint8ArraySource(bytes) {
-  return `new Uint8Array([${Array.from(bytes).join(',')}])`
-}
+const outPath = './src/models/index.ts'
 
 const [model, modelData, tokenizerModel] = await Promise.all([
   readFile(modelPath),
@@ -22,21 +19,40 @@ const [model, modelData, tokenizerModel] = await Promise.all([
 const ts = js`
 import * as ort from 'onnxruntime-web'
 import { SentencePieceProcessor } from '@agnai/sentencepiece-js'
+import { fromCompressed, fromBase64UrlString } from '@z-base/bytecodec'
 
 export async function createInferenceSession(): Promise<ort.InferenceSession> {
-  return ort.InferenceSession.create(${toUint8ArraySource(model)}, {
-    externalData: [
-      {
-        path: 'model.int4.onnx.data',
-        data: ${toUint8ArraySource(modelData)},
-      },
-    ],
-  })
+  return ort.InferenceSession.create(
+    await fromCompressed(
+      fromBase64UrlString(${JSON.stringify(toBase64UrlString(await toCompressed(model)))})
+    ),
+    {
+      externalData: [
+        {
+          path: 'model.int4.onnx.data',
+          data: await fromCompressed(
+            fromBase64UrlString(${JSON.stringify(toBase64UrlString(await toCompressed(modelData)))})
+          ),
+        },
+      ],
+    }
+  )
 }
 
-export async function createTokenProcessor():Promise<SentencePieceProcessor> {
+export async function createTokenProcessor(): Promise<SentencePieceProcessor> {
   const tokenProcessor = new SentencePieceProcessor()
-  await tokenProcessor.load(${toUint8ArraySource(tokenizerModel)})
+
+  const tokenizerModelBytes = await fromCompressed(
+    fromBase64UrlString(${JSON.stringify(toBase64UrlString(await toCompressed(tokenizerModel)))})
+  )
+
+  const tokenizerModelBlobBytes = Uint8Array.from(tokenizerModelBytes)
+
+  const tokenizerModelUrl = URL.createObjectURL(
+    new Blob([tokenizerModelBlobBytes], { type: 'application/octet-stream' })
+  )
+
+  await tokenProcessor.load(tokenizerModelUrl)
   return tokenProcessor
 }
 `.trimStart()

src/.types/index.d.ts

@@ -20,10 +20,17 @@ declare module '@agnai/sentencepiece-js' {
 
   export interface SentencePieceProcessorBinding {
     Load(model: SentencePieceStringViewHandle): SentencePieceStatus
-    EncodeAsIds(text: SentencePieceStringViewHandle): SentencePieceVector<number>
-    EncodeAsPieces(text: SentencePieceStringViewHandle): SentencePieceVector<string>
+    EncodeAsIds(
+      text: SentencePieceStringViewHandle
+    ): SentencePieceVector<number>
+    EncodeAsPieces(
+      text: SentencePieceStringViewHandle
+    ): SentencePieceVector<string>
     DecodeIds(ids: SentencePieceVector<number>): string
-    LoadVocabulary(vocab: SentencePieceStringViewHandle, threshold: number): SentencePieceStatus
+    LoadVocabulary(
+      vocab: SentencePieceStringViewHandle,
+      threshold: number
+    ): SentencePieceStatus
   }
 
   export interface SentencePieceModule {

src/Model/class.ts

@@ -1 +1 @@
-export { createInferenceSession, createTokenProcessor } from './Model/class.js'
+export { createInferenceSession, createTokenProcessor } from './models/index.js'

src/index.ts

@@ -0,0 +1,108 @@
+import * as ort from 'onnxruntime-web'
+import { cleanText } from '@agnai/sentencepiece-js'
+import { createInferenceSession, createTokenProcessor } from '../dist/index.js'
+
+const MAX_INPUT_LENGTH = 256
+const MAX_GENERATION_LENGTH = 513
+const BOS_TOKEN_ID = 1
+const EOS_TOKEN_ID = 2
+
+const input = cleanText(`- 3
+y TRADER JOE'S
+2001 Greenville Ave
+Dallas TX 75206
+Store #403 - (469) 334-0614
+OPEN 8:00AM TO 9:00PM DAILY
+R-CARROTS SHREDDED 10 0Z 1.29
+R-CUCUMBERS PERSIAN 1 LB 1.99
+TOMATOES CRUSHED NO SALT 1.59
+TOMATOES WHOLE NO SALT W/BASIL 1.59
+ORGANIC OLD_FASHIONED OATMEAL ~~ 2.69
+MINI-PEARL TOMATOES. . 2.49
+PKG SHREDDED MOZZARELLA LITET 3.9
+EGGS 1 DOZ ORGANIC BROWN. 3.79
+BEANS GARBANZO 0.89
+SPROUTED CA STYLE Zea
+A-AVOCADOS HASS BAG ACT 2:39
+A-APPLE BAG JAZZ 2 |B gr
+A-PEPPER BELL EACH XL RED 0.99
+GROCERY NON TAXABLE 0.98
+260.49
+BANANAS ORGANIC 0.87
+3kA 6 0.29/EA
+CREAMY SALTED PEANUT BUT TER 2.49
+WHL WHT PITA BREAD 1.69
+GROCERY NON TAXABLE 1.38
+260.69
+SUBTOTAL $38.68
+TOTAL $38.68
+CASH $40.00
+CHANGE $1.32
+ITEMS 22 Higgins, Ryan
+06-28-2014 12:34PM 0403 04 1346 4683
+THANK YOU FOR SHOPPING AT
+TRADER JOE'S
+www. trader joes .com
+`)
+
+function toInt64Tensor(values, dims) {
+  return new ort.Tensor('int64', BigInt64Array.from(values, BigInt), dims)
+}
+
+function argmax(values) {
+  let bestIndex = 0
+  let bestValue = Number.NEGATIVE_INFINITY
+
+  for (let index = 0; index < values.length; index += 1) {
+    if (values[index] > bestValue) {
+      bestValue = values[index]
+      bestIndex = index
+    }
+  }
+
+  return bestIndex
+}
+
+function getNextTokenId(logits) {
+  const [, targetLength, vocabSize] = logits.dims
+  const offset = (targetLength - 1) * vocabSize
+  const stepLogits = logits.data.subarray(offset, offset + vocabSize)
+
+  return argmax(stepLogits)
+}
+
+const tokenizer = await createTokenProcessor()
+const session = await createInferenceSession()
+
+const tokenIds = tokenizer.encodeIds(input).slice(0, MAX_INPUT_LENGTH)
+const attentionMask = tokenIds.map(() => 1)
+const decoderTokenIds = [BOS_TOKEN_ID]
+
+for (let step = 0; step < MAX_GENERATION_LENGTH; step += 1) {
+  const outputs = await session.run({
+    input_ids: toInt64Tensor(tokenIds, [1, tokenIds.length]),
+    attention_mask: toInt64Tensor(attentionMask, [1, attentionMask.length]),
+    decoder_input_ids: toInt64Tensor(decoderTokenIds, [
+      1,
+      decoderTokenIds.length,
+    ]),
+  })
+
+  const nextTokenId = getNextTokenId(outputs.logits)
+  if (nextTokenId === EOS_TOKEN_ID) {
+    break
+  }
+
+  decoderTokenIds.push(nextTokenId)
+}
+
+const outputTokenIds = decoderTokenIds.slice(1)
+const outputText = tokenizer.decodeIds(outputTokenIds)
+
+console.log({
+  inputLength: input.length,
+  tokenCount: tokenIds.length,
+  outputTokenCount: outputTokenIds.length,
+  outputTokenIds,
+  outputText,
+})