| from typing import Tuple | |
| import fasttext | |
| fasttext_model_path = "classifiers/ultra_fineweb_en.bin" | |
| # fasttext_model_path = "classifiers/ultra_fineweb_zh.bin" | |
| fasttext_model = fasttext.load_model(fasttext_model_path) | |
| def fasttext_infer(norm_content: str) -> Tuple[str, float]: | |
| """Fasttext inference function | |
| Args: | |
| content (str): input text | |
| Returns: | |
| str: json string with pred_label and pred_score | |
| """ | |
| pred_label, pred_prob = fasttext_model.predict(norm_content) | |
| pred_label = pred_label[0] | |
| _score = min(pred_prob.tolist()[0], 1) | |
| if pred_label == "__label__neg": | |
| _score = 1 - _score | |
| return pred_label, _score | |