Snapshot

completions.py

@@ -120,10 +120,10 @@ def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, de
     #%%
     words = split_into_words(token_probs, tokenizer)
     log_prob_threshold = -5.0
-    low_prob_words = [word for word in words if word.logprob < log_prob_threshold]
+    low_prob_words = [(i, word) for i, word in enumerate(words) if word.logprob < log_prob_threshold]
 
     #%%
-    contexts = [word.context for word in low_prob_words]
+    contexts = [word.context for _, word in low_prob_words]
     inputs = prepare_inputs(contexts, tokenizer, device)
     input_ids = inputs["input_ids"]
 
@@ -137,7 +137,12 @@ def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, de
     #%%
     replacements = extract_replacements(outputs, tokenizer, input_ids.shape[0], input_ids.shape[1], num_samples)
 
-    #%%
-    for word, replacements in zip(low_prob_words, replacements):
-        print(f"Original word: {word.text}, Log Probability: {word.logprob:.4f}")
-        print(f"Proposed replacements: {replacements}")
+    low_prob_words_with_replacements = { i: (w, r) for (i, w), r in zip(low_prob_words, replacements) }
+
+    result = []
+    for i, word in enumerate(words):
+        if i in low_prob_words_with_replacements:
+            result.append(ApiWord(text=word.text, logprob=word.logprob, replacements=low_prob_words_with_replacements[i][1]))
+        else:
+            result.append(ApiWord(text=word.text, logprob=word.logprob, replacements=[]))
+    return result

frontend/src/components/TokenChip.tsx

@@ -1,7 +1,5 @@
 import React, { useState } from "react"
 
-import React, { useState } from "react"
-
 export const TokenChip = ({
   token,
   logprob,

frontend/src/components/app.tsx

@@ -8,6 +8,13 @@ interface Word {
 }
 
 async function checkText(text: string): Promise<Word[]> {
+  const response = await fetch(`/check?text=${text}`)
+  const data = await response.json()
+  console.log(data)
+  return data.words
+}
+
+async function checkText0(text: string): Promise<Word[]> {
   await new Promise(resolve => setTimeout(resolve, 1000));
 
   const words = text.split(/\b/)

main.py

@@ -1,8 +1,7 @@
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel
 
-from models import CheckResponse, ApiWord
+from models import CheckResponse
 from completions import check_text, load_model
 
 app = FastAPI()
@@ -13,5 +12,4 @@ model, tokenizer, device = load_model()
 def check(text: str):
     return CheckResponse(text=text, words=check_text(text, model, tokenizer, device))
 
-# serve files from frontend/public
 app.mount("/", StaticFiles(directory="frontend/public", html=True))

models.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from pydantic import BaseModel
+
+@dataclass
+class Word:
+    tokens: list[int]
+    text: str
+    logprob: float
+    context: list[int]
+
+class ApiWord(BaseModel):
+    text: str
+    logprob: float
+    replacements: list[str]
+
+class CheckResponse(BaseModel):
+    text: str
+    words: list[ApiWord]