|
from __future__ import annotations |
|
from typing import NamedTuple |
|
import MeCab |
|
from transformers import PreTrainedTokenizer |
|
|
|
|
|
class MeCabResult(NamedTuple): |
|
"""MeCab解析結果の型 |
|
""" |
|
hyosokei: str |
|
hinshi: str |
|
|
|
|
|
class MeCabTokenizer(PreTrainedTokenizer): |
|
target_hinshi: list[str] | None |
|
mecab: MeCab.Tagger |
|
|
|
def __init__(self, |
|
hinshi: list[str] | None = None, |
|
mecab_dicdir: str | None = None, |
|
**kwargs): |
|
"""初期化処理 |
|
|
|
Args: |
|
hinshi (list[str] | None): 抽出する品詞 |
|
mecab_dicdir (str | None, optional): dicrcのあるディレクトリ |
|
""" |
|
|
|
self.target_hinshi = hinshi |
|
if mecab_dicdir is not None: |
|
self.mecab = MeCab.Tagger( |
|
f"-d {mecab_dicdir} -O '' -F '%m,%f[0]\n'") |
|
else: |
|
self.mecab = MeCab.Tagger("-O '' -F '%m,%f[0]\n'") |
|
|
|
super().__init__(**kwargs) |
|
|
|
def set_dicdir(self, mecab_dicdir: str): |
|
self.mecab = MeCab.Tagger(f"-d {mecab_dicdir} -O '' -F '%m,%f[0]\n'") |
|
|
|
def _tokenize(self, text: str) -> list[str]: |
|
"""文章から特定の品詞の単語を返します。 |
|
|
|
Args: |
|
text (str): 文章 |
|
|
|
Returns: |
|
list[str]: 特定の品詞の単語 |
|
""" |
|
|
|
out = [] |
|
|
|
result_words = self.mecab_analyze(text) |
|
for result_word in result_words: |
|
|
|
if result_word.hyosokei == "": |
|
continue |
|
if self.target_hinshi is not None: |
|
if result_word.hinshi in self.target_hinshi: |
|
|
|
out.append(result_word.hyosokei) |
|
else: |
|
continue |
|
else: |
|
out.append(result_word.hyosokei) |
|
return out |
|
|
|
def mecab_analyze(self, text: str) -> list[MeCabResult]: |
|
"""文章をMecabで分析します。 |
|
|
|
Args: |
|
text (str): 文章 |
|
|
|
Returns: |
|
list[MeCabResult]: MeCabの解析結果 |
|
""" |
|
nodes = self.mecab.parse(text).split("\n") |
|
out = [] |
|
for node in nodes: |
|
args = node.split(",") |
|
if args[0] in ["EOS", ""]: |
|
continue |
|
mecab_result = MeCabResult(args[0], args[1]) |
|
out.append(mecab_result) |
|
return out |
|
|
