wehub-resource-sync
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2022-09-28 13:31
|
|
import os
|
|
import sys
|
|
from typing import List, Union
|
|
|
|
import fasttext
|
|
from fasttext.FastText import _FastText
|
|
|
|
import hanlp
|
|
from hanlp.common.component import Component
|
|
from hanlp.utils.io_util import get_resource, stdout_redirected
|
|
from hanlp_common.io import load_json
|
|
from hanlp_common.reflection import classpath_of
|
|
from hanlp_common.structure import SerializableDict
|
|
|
|
|
|
class FastTextClassifier(Component):
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._model: _FastText = None
|
|
self.config = SerializableDict({
|
|
'classpath': classpath_of(self),
|
|
'hanlp_version': hanlp.__version__,
|
|
})
|
|
|
|
def load(self, save_dir, model_path=None, **kwargs):
|
|
config_path = os.path.join(save_dir, 'config.json')
|
|
if os.path.isfile(config_path):
|
|
self.config: dict = load_json(config_path)
|
|
model_path = self.config.get('model_path', model_path)
|
|
else:
|
|
model_path = model_path or save_dir
|
|
self.config['model_path'] = model_path
|
|
filepath = get_resource(model_path)
|
|
with stdout_redirected(to=os.devnull, stdout=sys.stderr):
|
|
self._model = fasttext.load_model(filepath)
|
|
|
|
def predict(self, text: Union[str, List[str]], topk=False, prob=False, max_len=None, **kwargs):
|
|
"""
|
|
Classify text.
|
|
|
|
Args:
|
|
text: A document or a list of documents.
|
|
topk: ``True`` or ``int`` to return the top-k labels.
|
|
prob: Return also probabilities.
|
|
max_len: Strip long document into ``max_len`` characters for faster prediction.
|
|
**kwargs: Not used
|
|
|
|
Returns:
|
|
Classification results.
|
|
"""
|
|
num_labels = len(self._model.get_labels())
|
|
flat = isinstance(text, str)
|
|
if flat:
|
|
text = [text]
|
|
if not isinstance(topk, list):
|
|
topk = [topk] * len(text)
|
|
if not isinstance(prob, list):
|
|
prob = [prob] * len(text)
|
|
if max_len:
|
|
text = [x[:max_len] for x in text]
|
|
text = [x.replace('\n', ' ') for x in text]
|
|
batch_labels, batch_probs = self._model.predict(text, k=num_labels)
|
|
results = []
|
|
for labels, probs, k, p in zip(batch_labels, batch_probs, topk, prob):
|
|
labels = [self._strip_prefix(x) for x in labels]
|
|
if k is False:
|
|
labels = labels[0]
|
|
elif k is True:
|
|
pass
|
|
elif k:
|
|
labels = labels[:k]
|
|
if p:
|
|
probs = probs.tolist()
|
|
if k is False:
|
|
result = labels, probs[0]
|
|
else:
|
|
result = dict(zip(labels, probs))
|
|
else:
|
|
result = labels
|
|
results.append(result)
|
|
if flat:
|
|
results = results[0]
|
|
return results
|
|
|
|
@property
|
|
def labels(self):
|
|
return [self._strip_prefix(x) for x in self._model.get_labels()]
|
|
|
|
@staticmethod
|
|
def _strip_prefix(label: str):
|
|
return label[len('__label__'):]
|