|
import spacy
|
|
|
|
__all__ = 'SpacyPipeline'
|
|
|
|
MODELS = {
|
|
"en": "en_core_web_sm",
|
|
"ja": "ja_core_news_sm",
|
|
"zh": "zh_core_web_sm",
|
|
"de": "de_core_news_sm",
|
|
"es": "es_core_news_sm",
|
|
"it": "it_core_news_sm",
|
|
"ko": "ko_core_news_sm",
|
|
"ru": "ru_core_news_sm",
|
|
"fr": "fr_core_news_sm",
|
|
"vi": "vi_core_news_lg"
|
|
}
|
|
VALID_METHODS = ['positionrank', 'textrank', 'biasedtextrank', 'positionrank', 'ner']
|
|
|
|
|
|
class SpacyPipeline:
|
|
|
|
def __init__(self, language, algorithm: str = None):
|
|
model = "vi_core_news_lg" if language not in MODELS else MODELS[language]
|
|
|
|
self.nlp = spacy.load(model)
|
|
self.nlp.add_pipe("sentencizer")
|
|
self.algorithm = algorithm
|
|
self.library = None
|
|
if self.algorithm is not None and self.algorithm != 'ner':
|
|
assert algorithm in VALID_METHODS, f'invalid algorithm {algorithm}\n- valid list: {VALID_METHODS}'
|
|
if self.algorithm == 'yake':
|
|
import spacy_ke
|
|
self.nlp.add_pipe("yake")
|
|
self.library = 'spacy_ke'
|
|
elif self.algorithm in ['textrank', 'biasedtextrank', 'positionrank']:
|
|
import pytextrank
|
|
self.nlp.add_pipe(algorithm)
|
|
self.library = 'pytextrank'
|
|
else:
|
|
raise ValueError(f'unknown algorithm: {self.algorithm}')
|
|
|
|
def _get_keyword(self, output, original_document=None, n=None):
|
|
if self.algorithm == 'ner':
|
|
return [str(i) for i in output.ents]
|
|
assert original_document is not None
|
|
assert n is not None
|
|
if self.library == 'spacy_ke':
|
|
return [str(term) for term, score in output._.extract_keywords(n) if str(term) in original_document]
|
|
return [str(i.text) for i in output._.phrases[:n] if str(i.text) in original_document]
|
|
|
|
def sentence_keyword(self, string: str, n: int = 10):
|
|
out = self.nlp(string)
|
|
sentence = [str(i) for i in out.sents if len(i) > 0]
|
|
keyword = self._get_keyword(out, string, n)
|
|
return sentence, keyword
|
|
|
|
def sentence(self, string: str):
|
|
return [str(i) for i in self.nlp(string).sents if len(i) > 0]
|
|
|
|
def token(self, string: str):
|
|
return [str(i) for i in self.nlp.tokenizer(string)]
|
|
|
|
def keyword(self, string: str, n: int = 10):
|
|
return self._get_keyword(self.nlp(string), string, n)
|
|
|
|
def ner(self, string: str, n: int = None):
|
|
keywords = self.nlp(string).ents
|
|
return keywords[:min(len(keywords), n)] if n is not None else keywords
|
|
|
|
@property
|
|
def language(self):
|
|
return self.nlp.lang |
