Improve get_sentences

src/great_ai/utilities/get_sentences.py

@@ -1,30 +1,42 @@
-from typing import List, cast
+import re
+from string import punctuation
+from typing import List
 
-from segtok.segmenter import split_multi
+from syntok.segmenter import segment
+from syntok.tokenizer import Tokenizer
 
 from .data import sentence_ending_punctuations
 
+punctuations_pattern = re.compile(f"\\s*[{re.escape(punctuation)}]+\\s*")
+
 
 def get_sentences(
-    text: str, ignore_partial: bool = False, true_case: bool = False
+    text: str,
+    ignore_partial: bool = False,
+    true_case: bool = False,
+    remove_punctuation: bool = False,
 ) -> List[str]:
-    if text.strip() == "":
-        return []
+    tokenizer = Tokenizer(
+        emit_hyphen_or_underscore_sep=True, replace_not_contraction=False
+    )
+    token_stream = tokenizer.tokenize(text)
 
-    possible_sentences = [
-        cast(str, s).strip() for s in split_multi(text) if cast(str, s).strip()
+    def process(sentence: str) -> str:
+        if true_case:
+            sentence = sentence[0].lower() + sentence[1:]  # very crude method
+        if remove_punctuation:
+            sentence = re.sub(punctuations_pattern, " ", sentence)
+        return sentence.strip()
+
+    sentences = [
+        process(tokenizer.to_text(sentence)) for sentence in segment(token_stream)
     ]
 
     if ignore_partial:
-        possible_sentences = [
-            s
-            for s in possible_sentences
-            if s[0].isupper() and s[-1] in sentence_ending_punctuations
+        sentences = [
+            sentence
+            for sentence in sentences
+            if sentence[0].isupper() and sentence[-1] in sentence_ending_punctuations
         ]
 
-    if true_case:
-        possible_sentences = [
-            s[0].lower() + s[1:] for s in possible_sentences  # very crude method
-        ]
-
-    return possible_sentences
+    return sentences