Update split_text.py

st_mlbee/split_text.py

@@ -0,0 +1,105 @@
+"""
+Split text to sentences.
+
+Modifed from seg_text.seg_text.py, vtext sentence_splitter removed, use
+
+Use sentence_splitter if supported,
+else use polyglot.text.Text
+
+!apt install libicu-dev
+!install pyicu pycld2
+!pip install polyglot sentence_splitter
+
+Use vtext and fastlid to rid of polyglot?
+
+from vtext.tokenize_sentence import UnicodeSentenceTokenizer, PunctuationTokenizer
+tok = UnicodeSentenceTokenizer()
+seg = tok.tokenize(''' Text ''') for langs not in LANG_S
+
+"""
+# pylint: disable=invalid-name
+
+import re
+from typing import List, Optional, Union
+
+import pysbd
+from fastlid import fastlid
+from loguru import logger
+from tqdm.auto import tqdm
+
+
+def _seg_text(
+    text: str,
+    lang: Optional[str] = None,
+) -> List[str]:
+    """
+    Split text to sentences.
+
+    Switched to pysbd
+
+    Args:
+    ----
+    text: string to split
+    lang: language, two-letter ISO (22 languages)
+
+    Returns:
+    -------
+    List of segmented sentences
+
+    """
+    if lang is None:
+        try:
+            lang, _ = fastlid(text)
+        except Exception as exc:
+            logger.warning(" fastlid: %s, setting lang='en'", exc)
+            lang = "en"
+
+    if not text.strip():
+        return []
+
+    seg = pysbd.Segmenter(language=lang, clean=True)
+
+    try:
+        # _ = tok.tokenize(text)
+        _ = seg.segment(text)
+    except Exception as exc:
+        logger.exception(f"pysbd.Segmenter, {exc=}")
+        raise
+    return _
+
+
+def seg_text(
+    lst: Union[str, List[str]],
+    lang: Optional[str] = None,
+    maxlines: int = 1000,
+    extra: Optional[str] = None,
+) -> List[str]:
+    """Split a list of text.
+
+    Arguments:
+        lst: text or text list
+        lang: optional lang code
+        maxlines: (default 1000), threshold for turn on tqdm progressbar, set to <1 or a large number to turn it off
+        extra: re.split(rf"{extra}, text) first
+    Returns:
+        list of splitted text.
+    """
+    if isinstance(lst, str):
+        lst = [lst]
+
+    if extra:
+        # insert \n
+        lst = [re.sub(rf"({extra})", r"\1\n", elm) for elm in lst]
+
+    res = []
+    for elm in lst:
+        res.extend(
+            _seg_text(
+                elm,
+                lang=lang,
+                maxlines=maxlines,
+                # flag=False,
+            )
+        )
+
+    return res