First working version

completions.py

@@ -91,13 +91,13 @@ def generate_outputs(model: PreTrainedModel, inputs: BatchEncoding, num_samples:
 def extract_replacements(outputs: GenerateOutput | torch.LongTensor, tokenizer: Tokenizer, num_inputs: int, input_len: int, num_samples: int = 5) -> list[list[str]]:
     all_new_words = []
     for i in range(num_inputs):
-        replacements = []
+        replacements = set()
         for j in range(num_samples):
             generated_ids = outputs[i * num_samples + j][input_len:]
             new_word = tokenizer.convert_ids_to_tokens(generated_ids.tolist())[0]
             if starts_with_space(new_word):
-                replacements.append(new_word[1:])
-        all_new_words.append(replacements)
+                replacements.add(" " +new_word[1:])
+        all_new_words.append(sorted(list(replacements)))
     return all_new_words
 
 #%%
@@ -128,7 +128,7 @@ def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, de
     input_ids = inputs["input_ids"]
 
     #%%
-    num_samples = 5
+    num_samples = 10
     start_time = time.time()
     outputs = generate_outputs(model, inputs, num_samples)
     end_time = time.time()

frontend/src/components/TokenChip.tsx

@@ -1,4 +1,4 @@
-import React, { useState } from "react"
+import React, { useState, useEffect, useRef } from "react"
 
 export const TokenChip = ({
   token,
@@ -14,6 +14,20 @@ export const TokenChip = ({
   onReplace: (newToken: string) => void
 }) => {
   const [isExpanded, setIsExpanded] = useState(false);
+  const dropdownRef = useRef<HTMLSelectElement>(null);
+
+  useEffect(() => {
+    const handleClickOutside = (event: MouseEvent) => {
+      if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) {
+        setIsExpanded(false);
+      }
+    };
+
+    document.addEventListener("mousedown", handleClickOutside);
+    return () => {
+      document.removeEventListener("mousedown", handleClickOutside);
+    };
+  }, []);
 
   const handleClick = () => {
     if (logprob < threshold && replacements.length > 0) {
@@ -39,6 +53,7 @@ export const TokenChip = ({
       {token}
       {isExpanded && (
         <select
+          ref={dropdownRef}
           onChange={handleReplacement}
           value={token}
           style={{

frontend/src/components/app.tsx

@@ -37,32 +37,40 @@ export default function App() {
   const [context, setContext] = useState("")
   const [wordlist, setWordlist] = useState("")
   const [showWholePrompt, setShowWholePrompt] = useState(false)
-  const [text, setText] = useState("The cumbersomely quick brown fox jumps over the conspicuously lazy dog.")
+  const [text, setText] = useState("I just drove to the store to but eggs, but they had some.")
   const [mode, setMode] = useState<"edit" | "check">("edit")
   const [words, setWords] = useState<Word[]>([])
   const [isLoading, setIsLoading] = useState(false)
 
+  const check = async () => {
+    setIsLoading(true)
+    try {
+      const checkedWords = await checkText(text)
+      setWords(checkedWords)
+    } finally {
+      setIsLoading(false)
+      setMode("check")
+    }
+  }
+
   const toggleMode = async () => {
     if (mode === "edit") {
       setIsLoading(true)
-      try {
-        const checkedWords = await checkText(text)
-        setWords(checkedWords)
-      } finally {
-        setMode("check")
-        setIsLoading(false)
-      }
+      await check()
     } else {
       setMode("edit")
     }
   }
 
-  const handleReplace = (index: number, newToken: string) => {
+  const handleReplace = async (index: number, newToken: string) => {
     const updatedWords = [...words]
     updatedWords[index].text = newToken
+    updatedWords[index].logprob = 0
+    updatedWords[index].replacements = []
     setWords(updatedWords)
     setText(updatedWords.map(w => w.text).join(""))
-    setMode("edit")
+    // setMode("edit")
+    await check()
   }
 
   let result
@@ -101,7 +109,7 @@ export default function App() {
       <details>
         <summary>Advanced settings</summary>
         <label>
-          <strong>Threshold:</strong> <input type="number" step="0.1" value={threshold} onChange={e => setThreshold(Number(e.target.value))} />
+          <strong>Threshold:</strong> <input type="number" step="1" value={threshold} onChange={e => setThreshold(Number(e.target.value))} />
           <small>
             The <a href="https://en.wikipedia.org/wiki/Log_probability" target="_blank" rel="noreferrer">logprob</a> threshold.
             Tokens with logprobs smaller than this will be marked red.
@@ -132,8 +140,8 @@ export default function App() {
 
         <p>
           <small>
-            Made by <a href="https://vgel.me">Theia Vogel</a> (<a href="https://twitter.com/voooooogel">@voooooogel</a>).
-            Made with Svelte, GPT-3, and transitively, most of the web.
+            Based on <a href="https://github.com/vgel/gpted">GPTed</a> by <a href="https://vgel.me">Theia Vogel</a>.
+            Made with React, Transformers, LLama 3.2, and transitively, most of the web.
             <br />
             This software is provided with absolutely no warranty.
           </small>