demo_detoxi / model_wrapper /fasttext_wrapper.py
Cricles's picture
Update model_wrapper/fasttext_wrapper.py
0d3b9b0 verified
raw
history blame
284 Bytes
import fasttext
class FasttextWrapper(object):
MODEL_PATH: str = "./model_fasttext.bin"
def __init__(self) -> None:
self.model = fasttext.load_model(FasttextWrapper.MODEL_PATH)
def __call__(self, text: str) -> str:
return self.model.predict(text)[0][0]