File size: 919 Bytes
06bc836 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | from fastapi import FastAPI
from pydantic import BaseModel
import kenlm
app = FastAPI()
lm = kenlm.Model("char.bin")
CORPUS = open(
"1.txt",
encoding="utf-8"
).read().splitlines()
class Query(BaseModel):
text: str
def generate_candidates(prefix, max_n=100):
cands = []
for line in CORPUS:
words = line.split()
for w in words:
if w.startswith(prefix):
cands.append(w)
return list(set(cands))[:max_n]
@app.post("/predict")
def predict(q: Query):
prefix = q.text.strip()
candidates = generate_candidates(prefix)
scored = []
for c in candidates:
score = lm.score(c)
scored.append({
"word": c,
"score": score
})
scored.sort(key=lambda x: x["score"], reverse=True)
return {
"candidates": scored[:5]
}
|