Update index.html

index.html

@@ -3,84 +3,50 @@
 <head>
   <meta charset="UTF-8">
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Swedish → English Translator</title>
 
-  <title>English → Swedish Translator</title>
-
-  <script
-    type="module"
-    crossorigin
-    src="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.js">
-  </script>
-
-  <link
-    rel="stylesheet"
-    crossorigin
-    href="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.css">
+  <script src="https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js"></script>
 
   <style>
-    body {
-      max-width: 900px;
-      margin: 0 auto;
-      padding: 24px;
-      font-family: Arial, sans-serif;
-    }
-
-    #status {
-      color: #555;
-      margin: 8px 0 18px;
-    }
-
-    #error {
-      color: #b00020;
-      white-space: pre-wrap;
-      margin-top: 12px;
-    }
-
-    textarea {
-      width: 100%;
-      min-height: 120px;
-      padding: 12px;
-      box-sizing: border-box;
-      font-size: 16px;
-    }
-
-    button {
-      margin: 14px 0;
-      padding: 10px 18px;
-      font-size: 16px;
-      cursor: pointer;
-    }
+    body { max-width: 900px; margin: 0 auto; padding: 24px; font-family: Arial, sans-serif; }
+    #status { color: #555; margin: 8px 0 18px; }
+    #error { color: #b00020; white-space: pre-wrap; margin-top: 12px; }
+    textarea { width: 100%; min-height: 120px; padding: 12px; box-sizing: border-box; font-size: 16px; }
+    button { margin: 14px 0; padding: 10px 18px; font-size: 16px; cursor: pointer; }
   </style>
 </head>
 
 <body>
-  <h1>English → Swedish Translator</h1>
+  <h1>Swedish → English Translator</h1>
   <p>
-    This translation model runs directly in your browser.
-    The first translation can take a little while because the model must download.
+    This is our own character-level Transformer from the lab, running directly in your browser.
+    Sentences up to 120 characters work best (that is what it was trained on).
   </p>
 
-  <p id="status">Loading translation model…</p>
+  <p id="status">Loading model…</p>
 
-  <label for="input">English text</label>
-  <textarea
-    id="input"
-    placeholder="Write an English sentence here..."
-    disabled
-  ></textarea>
+  <label for="input">Swedish text</label>
+  <textarea id="input" placeholder="Skriv en svensk mening här..." disabled></textarea>
 
   <br>
-
   <button id="translate-button" disabled>Translate</button>
 
-  <label for="output">Swedish translation</label>
+  <label for="output">English translation</label>
   <textarea id="output" readonly></textarea>
 
   <p id="error"></p>
 
-  <script type="module">
-    import { pipeline, env } from
-      "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.2";
+  <script>
+    // Files produced by export_onnx.py, placed next to this page
+    const ENCODER_URL = "encoder.onnx";
+    const DECODER_URL = "decoder.onnx";
+    const VOCAB_URL   = "vocab.json";
+
+    // Same constants as the notebook
+    const TRAIN_MAX_LEN = 120;              // max_len
+    const MAX_LEN = TRAIN_MAX_LEN + 2;      // + <s> and </s>
+
+    ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/";
 
     const status = document.getElementById("status");
     const input = document.getElementById("input");
@@ -88,32 +54,68 @@
     const button = document.getElementById("translate-button");
     const errorBox = document.getElementById("error");
 
-    let translator;
+    let encoder, decoder, token2idx, idx2token;
 
     async function loadModel() {
       try {
-        env.allowLocalModels = false;
+        token2idx = await (await fetch(VOCAB_URL)).json();
+        idx2token = Object.fromEntries(Object.entries(token2idx).map(([t, i]) => [i, t]));
 
-        translator = await pipeline(
-          "translation",
-          "Xenova/opus-mt-en-sv"
-        );
+        const opts = { executionProviders: ["wasm"] };
+        encoder = await ort.InferenceSession.create(ENCODER_URL, opts);
+        decoder = await ort.InferenceSession.create(DECODER_URL, opts);
 
-        status.textContent = "Model ready. Enter English text and click Translate.";
+        status.textContent = "Model ready. Enter Swedish text and click Translate.";
         input.disabled = false;
         button.disabled = false;
       } catch (error) {
         console.error(error);
-        status.textContent = "The translation model could not load.";
+        status.textContent = "The model could not load.";
         errorBox.textContent = String(error);
       }
     }
 
-    button.addEventListener("click", async () => {
-      const text = input.value.trim();
+    // seqs2tensors: <s> + characters + </s>, unknown characters -> <unk> (1)
+    function encode(text) {
+      const chars = Array.from(text);        // iterates by code point, like Python's list(str)
+      return [token2idx["<s>"], ...chars.map(c => token2idx[c] ?? 1), token2idx["</s>"]];
+    }
+
+    const toTensor = ids =>
+      new ort.Tensor("int64", BigInt64Array.from(ids, BigInt), [1, ids.length]);
+
+    function argmax(arr) {
+      let best = 0;
+      for (let i = 1; i < arr.length; i++) if (arr[i] > arr[best]) best = i;
+      return best;
+    }
+
+    // Same loop as greedy_decode_batched in the notebook
+    async function translate(text) {
+      const src = encode(text);
+      const memory = (await encoder.run({ src: toTensor(src) })).memory;
+
+      const eos = token2idx["</s>"];
+      const tgt = [token2idx["<s>"]];
+      const maxLen = Math.min(src.length + 20, MAX_LEN);
 
+      for (let step = 0; step < maxLen - 1; step++) {
+        const logits = (await decoder.run({ tgt: toTensor(tgt), memory })).logits.data;
+        const next = argmax(logits);
+        tgt.push(next);
+        if (next === eos) break;
+      }
+
+      return tgt.slice(1)
+        .filter(i => i !== eos)
+        .map(i => idx2token[i] ?? "<unk>")
+        .join("");
+    }
+
+    async function run() {
+      const text = input.value.trim();
       if (!text) {
-        output.value = "Please write something first.";
+        output.value = "Skriv något först.";
         return;
       }
 
@@ -122,9 +124,10 @@
       errorBox.textContent = "";
 
       try {
-        const result = await translator(text);
-        output.value = result[0].translation_text;
-        status.textContent = "Done.";
+        output.value = await translate(text);
+        status.textContent = text.length > TRAIN_MAX_LEN
+          ? `Done. Note: the input is longer than ${TRAIN_MAX_LEN} characters, so the result may be poor.`
+          : "Done.";
       } catch (error) {
         console.error(error);
         status.textContent = "Translation failed.";
@@ -132,9 +135,14 @@
       } finally {
         button.disabled = false;
       }
+    }
+
+    button.addEventListener("click", run);
+    input.addEventListener("keydown", e => {
+      if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); run(); }
     });
 
     loadModel();
   </script>
 </body>
-</html>
+</html>