File size: 2,485 Bytes
97c46f0 cd708ec 97c46f0 ef06cf1 97c46f0 ef06cf1 97c46f0 ef06cf1 97c46f0 301bdd5 97c46f0 301bdd5 ef06cf1 97c46f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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 = []
# Mecabで分析します。
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
|