Merge pull request #22 from DoodleBears/21-fix-bad-long-text-split

split_lang/split/splitter.py

@@ -29,7 +29,9 @@ def __init__(
         default_lang: str = DEFAULT_LANG,
         punctuation: str = PUNCTUATION,
         not_merge_punctuation: str = "\n",
+        special_merge_for_zh_ja: bool = True,
         merge_across_punctuation: bool = True,
+        merge_across_newline: bool = True,
         merge_across_digit: bool = True,
         debug: bool = False,
     ) -> None:
@@ -52,6 +54,8 @@ def __init__(
         self.debug = debug
         self.punctuation = punctuation
         self.not_merge_punctuation = not_merge_punctuation
+        self.merge_across_newline = merge_across_newline
+        self.special_merge_for_zh_ja = special_merge_for_zh_ja
         self.merge_across_punctuation = merge_across_punctuation
         self.merge_across_digit = merge_across_digit
 
@@ -78,13 +82,24 @@ def split_by_lang(
             substrings.extend(section.substrings)
         substrings = self._merge_digit(substrings=substrings)
 
+        if self.merge_across_newline:
+            substrings = self._merge_substrings_across_newline(
+                substrings=substrings,
+            )
         if self.merge_across_digit:
             substrings = self._merge_substring_across_digit(substrings=substrings)
 
         if self.merge_across_punctuation:
             substrings = self._merge_substrings_across_punctuation(
                 substrings=substrings,
             )
+        if self.merge_across_newline:
+            substrings = self._merge_substrings_across_newline(
+                substrings=substrings,
+            )
+        if self.special_merge_for_zh_ja:
+            substrings = self._special_merge_for_zh_ja(substrings=substrings)
+            substrings = self._special_merge_for_zh_ja(substrings=substrings)
 
         return substrings
 
@@ -586,38 +601,125 @@ def _merge_digit(
         new_substrings = self._merge_substrings(substrings=substrings)
         return new_substrings
 
-    # MARK: _merge_substring_across_digit
-    def _merge_substring_across_digit(
+    def _merge_substrings_across_newline(
         self,
         substrings: List[SubString],
     ) -> List[SubString]:
         new_substrings: List[SubString] = []
+        last_lang = ""
+        for index in range(len(substrings)):
+            if (
+                last_lang == substrings[index].lang
+                or substrings[index].lang == "newline"
+            ):
+                new_substrings[-1].text += substrings[index].text
+                new_substrings[-1].length += substrings[index].length
+            else:
+                new_substrings.append(substrings[index])
+                last_lang = substrings[index].lang
+        return new_substrings
 
-        for _, substring in enumerate(substrings):
-            if substring.is_digit:
+    # MARK: _special_merge_for_zh_ja
+
+    def _special_merge_for_zh_ja(
+        self,
+        substrings: List[SubString],
+    ) -> List[SubString]:
+        new_substrings: List[SubString] = []
+
+        if len(substrings) == 1:
+            return substrings
+        # TODO: 如果遇到两边的语言相同，即 ABA，且两边的 A 的长度合计加起来大于中间的 B 的长度的10倍，则合并
+
+        for index in range(len(substrings)):
+            current_block = substrings[index]
+            if index == 0:
                 if (
-                    new_substrings
-                    and new_substrings[-1].lang != "punctuation"
-                    and new_substrings[-1].lang != "newline"
+                    substrings[index + 1].lang in ["zh", "ja"]
+                    and substrings[index].lang in ["zh", "ja", "x"]
+                    and substrings[index].length * 10 < substrings[index + 1].length
                 ):
-                    new_substrings[-1].text += substring.text
-                    new_substrings[-1].length += substring.length
+                    right_block = substrings[index + 1]
+                    new_substrings.append(
+                        SubString(
+                            is_digit=False,
+                            is_punctuation=False,
+                            lang=right_block.lang,
+                            text=current_block.text + right_block.text,
+                            length=current_block.length + right_block.length,
+                            index=current_block.index,
+                        )
+                    )
+                    index += 1
                 else:
-                    new_substrings.append(substring)
+                    new_substrings.append(substrings[index])
+            elif index == len(substrings) - 1:
+                left_block = new_substrings[-1]
+                if (
+                    left_block.lang in ["zh", "ja"]
+                    and current_block.lang in ["zh", "ja", "x"]
+                    and current_block.length * 10 < left_block.length
+                ):
+                    new_substrings[-1].text += current_block.text
+                    new_substrings[-1].length += current_block.length
+
+                    index += 1
+                else:
+                    new_substrings.append(substrings[index])
             else:
                 if (
-                    new_substrings
-                    and new_substrings[-1].lang == "digit"
-                    and substring.lang != "newline"
+                    new_substrings[-1].lang == substrings[index + 1].lang
+                    and new_substrings[-1].lang in ["zh", "ja"]
+                    # and substrings[index].lang in ["zh", "ja", "x"]
+                    and substrings[index].lang != "en"
+                    and substrings[index].length * 10
+                    < new_substrings[-1].length + substrings[index + 1].length
                 ):
-                    temp = new_substrings.pop()
-                    temp.text = temp.text + substring.text
-                    temp.index = temp.index
-                    temp.length = temp.length + substring.length
-                    substring = temp
+                    left_block = new_substrings[-1]
+                    right_block = substrings[index + 1]
+                    current_block = substrings[index]
+                    new_substrings[-1].text += current_block.text + right_block.text
+                    new_substrings[-1].length += (
+                        current_block.length + right_block.length
+                    )
+                    index += 1
+                else:
+                    new_substrings.append(substrings[index])
+        new_substrings = self._merge_substrings(substrings=new_substrings)
+        return new_substrings
+
+    # MARK: _merge_substring_across_digit
+    def _merge_substring_across_digit(
+        self,
+        substrings: List[SubString],
+    ) -> List[SubString]:
+        new_substrings: List[SubString] = []
+        last_lang = ""
+
+        for _, substring in enumerate(substrings):
+            if new_substrings:
+                if substring.is_digit:
+                    if new_substrings[-1].is_digit:
+                        new_substrings[-1].text += substring.text
+                        new_substrings[-1].length += substring.length
+                    else:
+                        new_substrings[-1].text += substring.text
+                        new_substrings[-1].length += substring.length
+                else:
+                    if substring.lang == last_lang or last_lang == "":
+                        new_substrings[-1].text += substring.text
+                        new_substrings[-1].length += substring.length
+                        new_substrings[-1].lang = (
+                            substring.lang
+                            if new_substrings[-1].lang == "digit"
+                            else new_substrings[-1].lang
+                        )
+                    else:
+                        new_substrings.append(substring)
+                    last_lang = substring.lang
+            else:
                 new_substrings.append(substring)
 
-        new_substrings = self._merge_substrings(substrings=new_substrings)
         return new_substrings
 
     # MARK: _merge_substrings_across_punctuation

tests/data/correct_split.txt

@@ -26,4 +26,5 @@ Bonjour|, |wie geht's dir |today|?
 Vielen Dank |merci beaucoup |for your help|.
 Ich bin müde |je suis fatigué |and I need some rest|.
 Ich mag dieses Buch |ce livre est intéressant |and it has a great story|.
-Ich mag dieses Buch|, |ce livre est intéressant|, |and it has a great story|.
+Ich mag dieses Buch|, |ce livre est intéressant|, |and it has a great story|.
+入学早々隣のクラスの絵奈は男子からの人気者。新入生代表の挨拶も絵奈だったから目立ってしまって余計にあの子誰？ってなった。絵奈と私には近所に住む幼稚園の頃から仲の良い幼なじみの澤村　和樹（さわむら　かずき）という同じ中学の同級生でもある男子生徒がいる。絵奈も私もかずくんと呼んでいる。かずくんは小学１年生の頃からサッカーをやっていて、高学年になる頃には県の選抜に選ばれたりと、丸山台中の新入生では期待の星。何でユースとかにいかないで、公立の中学校でサッカー部なのかはわからない。かずくんとは絵奈より私の方が話が合うことが多かった。小学生の頃もよく３人で遊んだりもしたし、中学生になってからも３人で朝一緒に行こうなんて約束していた。迎えた入学式後の初登校日。「忘れ物とかない？気を付けて行くんだよ。急がなくても充分間に合うからね。」私にだけそう言ったママ。絵奈には「いってらっしゃい」だけ。（二人とも一緒に行くのに、私だけ落ち着きないみたいに思われてる…？）「ほら、早く行くよ。ママ、いってきます♪」私の横で姉の絵奈が私に向かってそう言うと、ママには笑顔で挨拶をしていた。そそくさと出てしまう絵奈。「ちょ、ちょっと（汗）充分間に合うからねってママ言ってるじゃん（汗）んじゃいってきま〜す！」そう言って、私も絵奈の後を急いで追う。外に出るとかずくんがちょっと遠くから歩いてきているのが見えたので、私は大声で手を振って呼んだ。「かずく〜ん！！おはよー！！」私の大声に気づいたかずくんは走ってきた。「よぉ！！絵奈と残念な方の妹！（笑）」（こいつ…面白がってやがるな…）絵奈が新入生代表で目立ってしまって、私は自分のクラスで残念な方の妹だと、昨日、ふざけた男子に言われたのだ…そしてかずくんと私は同じクラスなのでそれを面白がって言ってきたのだろう…うるせぇーよって言ってやろうと思ったら、私より先に絵奈が口を開いた。「かずくん、残念な方の妹って何？」絵奈に聞かれるとちょっとバツが悪そうなかずくんが渋々答える。「いや、その、実は俺達のクラスで、ふざけて紗羅が残念な方の妹って男子にからかわれてたんだよ（汗）」「は？で、かずくんも面白がって、そんなこと言ってるわけ？くだらないことやめなよ。そういう時は普通、助ける男がカッコいいんでしょ。一緒になって紗羅が気にしてないからって最悪だよ。」絵奈はどちらかと言うと温厚で優しいタイプなのに少し怒っているように言ったので、かずくんはしょんぼり、私も珍しく絵奈が怒っているので少しびっくりした…３人でちょっと気まずくなって歩くのに耐えられない私が絵奈に言った。「まぁまぁ、そんなツンケンしないでさ（汗）絵奈らしくないよ？かずくんも悪気があって言ったわけじゃないんだし、許してあげなよ（笑）」「私らしいって何？紗羅は悔しくないの？残念な方の妹とか言われて。そもそも紗羅がそんなお気楽な感じだから駄目なんでしょ（怒）私は紗羅が悪口言われてたら自分のことのように面白くないよ（怒）」（怒りの矛先が私に向き始めたな…）絵奈は温厚で優しく優等生タイプだけど、怒るとちょっと面倒くさいタイプ…何日も根に持つし、頑固なところもある。私は対照的でカッとなっても、１日経てばそんなこと忘れてるし、あんまり根に持つこともない。頑固な所は同じかもしれないけど。怒り出すと絵奈をなだめるのは結構大変…普段あんまり怒らないタイプだから、かもしれないけど…「そのさ、絵奈は私と違っておとなしめな品のあるタイプなんだから、あんまりカッカしてるとイメージ崩れるよ？なんちゃって（笑）かずくんも少しびっくりしてるでしょ（汗）ほらほら、せっかく入学して登校初日なのに空気悪いのおしまいにして楽しく行こうよ♪（笑）」一応、私がそう言うとカッカしてた絵奈も機嫌を少し直してくれて、かずくんもごめんって謝って一件落着し、その後は仲良く登校し中学校へ無事に到着。「んじゃまた帰りね〜♪」そう言って、絵奈とは別々の教室へ。私はかずくんと同じクラスへ入る。先生が来て、授業が始まる。勉強は嫌いだ…（今日の給食なんだっけかなぁ〜）なんて授業に集中していないと見事に先生に当てられているのに気づかない。隣の席のかずくんが小さな声で「おーい。紗羅。おーい」って言っているのに気づく。（ん？何？教科書忘れたのかこいつは本当にサッカーしかできない奴だからなぁ。しょうがないなぁ）机をくっつけて教科書を見せてやろうとするとかずくんはオドオドする。「ちげーよ（汗）紗羅（汗）」また小さな声でかずくんが言った後には怒った先生が私の目の前にいた…「山本、お前何やってんだ？あくびしたり、外をボケーっと見てたり。同じ山本のお前の姉とは大違いだな。」先生は不機嫌そうに言う。（すいませんね〜絵奈だったら、ちゃんと授業聞いてるんでしょうね〜私は残念な方の妹なので（笑））なんて怒られているはずなのに思っていると、朝、絵奈から怒られたからか、かずくんが私を庇った。「その、先生。絵奈は絵奈。紗羅は紗羅ですよ。」かずくんの言葉に先生もちょっと悪いと思ったのか、ブツブツ言いながら戻っていき、私がボケーっとしていたことは流されて終わった。（おぉーかずくんが助けてくれた（笑）もしかして、絵奈に朝怒られたから？もしかして、かずくんも昔から絵奈のことが…いや、考えるのはやめよう。私達３人は幼馴染。変に関係がおかしくなるのも嫌だし。）なんてふと思った。そして、この先、私達３人の関係は壊れていくことにこの時は気づかなかった…