Rekipjan's picture
Upload folder using huggingface_hub
06bc836 verified
raw
history blame contribute delete
919 Bytes
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]
}