Cricles commited on
Commit
a696648
·
verified ·
1 Parent(s): edf196f

Update model_wrapper/fasttext_wrapper.py

Browse files
Files changed (1) hide show
  1. model_wrapper/fasttext_wrapper.py +2 -1
model_wrapper/fasttext_wrapper.py CHANGED
@@ -7,4 +7,5 @@ class FasttextWrapper(object):
7
  self.model = fasttext.load_model(FasttextWrapper.MODEL_PATH)
8
 
9
  def __call__(self, text: str) -> str:
10
- return self.model.predict(text)[0][0]
 
 
7
  self.model = fasttext.load_model(FasttextWrapper.MODEL_PATH)
8
 
9
  def __call__(self, text: str) -> str:
10
+ label, prob=self.model.predict(text)
11
+ return label[0],prob