commit files to HF hub

mecab_tokenizer.py

@@ -9,14 +9,6 @@ class MeCabResult(NamedTuple):
     """
     hyosokei: str
     hinshi: str
-    hinshi_saibunrui_1: str
-    hinshi_saibunrui_2: str
-    hinshi_saibunrui_3: str
-    katsuyokei_1: str
-    katsuyokei_2: str
-    genkei: str = ""
-    yomi: str = ""
-    hatsuon: str = ""
 
 
 class MeCabTokenizer(PreTrainedTokenizer):
@@ -34,12 +26,16 @@ class MeCabTokenizer(PreTrainedTokenizer):
 
         self.target_hinshi = hinshi
         if mecab_dicdir is not None:
-            self.mecab = MeCab.Tagger(f"-d {mecab_dicdir} -O '' -F '%m,%H\n'")
+            self.mecab = MeCab.Tagger(
+                f"-d {mecab_dicdir} -O '' -F '%m,%f[0]\n'")
         else:
-            self.mecab = MeCab.Tagger("-O '' -F '%m,%H\n'")
+            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]:
         """文章から特定の品詞の単語を返します。
 
@@ -77,29 +73,11 @@ class MeCabTokenizer(PreTrainedTokenizer):
             list[MeCabResult]: MeCabの解析結果
         """
         nodes = self.mecab.parse(text).split("\n")
-        #形態素1つ1つを処理
         out = []
         for node in nodes:
             args = node.split(",")
             if args[0] in ["EOS", ""]:
                 continue
-            # 辞書によって異なる
-            if len(args) == 10:
-                mecab_result = MeCabResult(args[0], args[1], args[2], args[3],
-                                           args[4], args[5], args[6], args[7],
-                                           args[8], args[9])
-            elif len(args) == 7:
-                # 英語
-                mecab_result = MeCabResult(args[0], args[1], args[2], args[3],
-                                           args[4], args[5], args[6], "*", "*",
-                                           "*")
-            elif len(args) == 27:
-                # 補助記号
-                mecab_result = MeCabResult(args[0], args[1], args[2], args[3],
-                                           args[4], args[5], args[6], "*", "*",
-                                           "*")
-            else:
-                raise NotImplementedError(
-                    f"unsupposed parse args_length: {len(args)}, args: {args}")
+            mecab_result = MeCabResult(args[0], args[1])
             out.append(mecab_result)
         return out