Cleanup

completions.py

@@ -71,23 +71,6 @@ def calculate_log_probabilities(model: PreTrainedModel, tokenizer: Tokenizer, in
     tokens: torch.Tensor = input_ids[0][1:]
     return list(zip(tokens.tolist(), token_log_probs.tolist()))
 
-def generate_outputs(model: PreTrainedModel, inputs: BatchEncoding, num_samples: int = 5) -> GenerateOutput | torch.LongTensor:
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=4,
-            num_return_sequences=num_samples,
-            temperature=1.0,
-            top_k=50,
-            top_p=0.95,
-            do_sample=True
-            # num_beams=num_samples
-        )
-    return outputs
-
 #%%
 
 def load_model() -> tuple[PreTrainedModel, Tokenizer, torch.device]:

main.py

@@ -2,29 +2,13 @@ from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from functools import lru_cache
 
-from models import ApiWord, CheckResponse
+from models import CheckResponse
 from completions import check_text, load_model
 
 app = FastAPI()
 
 model, tokenizer, device = load_model()
 
-def check_text_stub(text: str):
-    def rep(i):
-        if i == 3:
-            return -10, [" jumped", " leaps"]
-        if i == 5:
-            return -10, [" calm"]
-        if i == 7:
-            return -10, [" dog", " cat", " bird", " fish"]
-        return -3, []
-
-    result = []
-    for i, w in enumerate(text.split()):
-        logprob, replacements = rep(i)
-        result.append(ApiWord(text=f" {w}", logprob=logprob, replacements=replacements))
-    return result
-
 @lru_cache(maxsize=100)
 def cached_check_text(text: str):
     return check_text(text, model, tokenizer, device)
@@ -34,5 +18,3 @@ def check(text: str):
     return CheckResponse(text=text, words=cached_check_text(text))
 
 app.mount("/", StaticFiles(directory="frontend/public", html=True))
-
-#%%