Add option to TextSplitter to return individual sentences. Adding general SaT model support.

[hhuangMITRE]

Issues:

• Add option to TextSplitter to return individual sentences openmpf#1965

Related PRs:

• Add option to TextSplitter to return individual sentences. Adding general SaT model support. openmpf-components#408

Summary:

This PR updates the nlp_text_splitter to add SaT (segment any text, https://github.com/segment-any-text/wtpsplit) model support, newline handling, sentence splitting options, and preferred-limit chunking.

Before this PR, only WtP and spaCy models are available for text segmentation. There also was a general segmentation strategy: estimate a sentence break near a hard size limit, then walk back to the nearest sentence boundary to generate the largest possible chunk. This works well for components with a large character or token text limit, however it may create issues for other components where a smaller text limit is needed (where possible).

This update adds:

• SaT model support.

• Newline normalization so line breaks can be treated as spaces, removed, or preserved depending on language/script estimated.

• New options for chunk vs individual sentence splitting.

• A new preferred/soft limit so users can try to generate smaller chunks while still respecting a hard text limit.

• Along the way, an improved breakpoint alignment function was added so sentence boundaries are computed against the original text, as SaT and WtP appeared to remove extraneous whitespace while splitting. Additional changes were made to handle edge cases for empty outputs, zero-length text, and mid-word splits

• Finally, support for Flores/NLLB language codes in WtpLanguageSettings is being transferred over to this PR.

---

This change is 

[jrobble]

@jrobble reviewed 4 of 4 files at r1, all commit messages.
Reviewable status: all files reviewed, 4 unresolved discussions (waiting on @hhuangMITRE)

---

a discussion (no related file):
Mention SaT here:

# To hold spaCy, WtP, and other potential sentence detection models in cache

Mention SaT here:

            log.warning(
                "Invalid model setting '%s'. Only `cpu` and `cuda` "
                        "(or `gpu`) WtP model options available at this time. "
                        "Defaulting to `cpu` mode.", model_setting)

Mention SaT in install.sh and LICENSE.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 83 at r1 (raw file):

            self._update_wtp_model(model_name, model_setting, default_lang)
            self.split = self._split_wtp
            log.info("Setup WtP model: %s", model_name)

Generally, 'f' strings are preferred since they keep the variable name inline with the text. It makes things easier to read.

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 68 at r1 (raw file):

        self.assertEqual(2, len(actual))
        self.assertEqual('Hello, what is your name? ', actual[0])
        self.assertEqual('My name is John.', actual[1])

These asserts as the same as above test_sat_basic_sentence_split test. I would feel better if we can prove that the different splitting behaviors return different results.

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 104 at r1 (raw file):

            500,
            len,
            self.sat_model,split_mode=SplitMode.SENTENCE))

Formatting nitpick: Move split_mode to next line.

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 106 at r1 (raw file):

            self.sat_model,split_mode=SplitMode.SENTENCE))
        self.assertEqual(input_text, ''.join(actual))
        self.assertEqual(2, len(actual))

These asserts as the same as above. I would feel better if we can prove that the different splitting behaviors return different results.

[hhuangMITRE]

Reviewable status: 0 of 8 files reviewed, 4 unresolved discussions (waiting on @hhuangMITRE and @jrobble)

---

a discussion (no related file):

Previously, jrobble (Jeff Robble) wrote…

Mention SaT here:

# To hold spaCy, WtP, and other potential sentence detection models in cache

Mention SaT here:

            log.warning(
                "Invalid model setting '%s'. Only `cpu` and `cuda` "
                        "(or `gpu`) WtP model options available at this time. "
                        "Defaulting to `cpu` mode.", model_setting)

Mention SaT in install.sh and LICENSE.

Done.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 83 at r1 (raw file):

Previously, jrobble (Jeff Robble) wrote…

Generally, 'f' strings are preferred since they keep the variable name inline with the text. It makes things easier to read.

Updated, thanks!

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 68 at r1 (raw file):

Previously, jrobble (Jeff Robble) wrote…

These asserts as the same as above test_sat_basic_sentence_split test. I would feel better if we can prove that the different splitting behaviors return different results.

I've added in the new test cases. There's also some new differences in translation which I've added to the other PR.

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 104 at r1 (raw file):

Previously, jrobble (Jeff Robble) wrote…

Formatting nitpick: Move split_mode to next line.

Done!

---

detection/nlp_text_splitter/tests/test_text_splitter.py line 106 at r1 (raw file):

Previously, jrobble (Jeff Robble) wrote…

These asserts as the same as above. I would feel better if we can prove that the different splitting behaviors return different results.

I've tweaked the test, right now SaT seems more sensitive to splitting it seems.

[gonzalezjo]

Edit: this is my first time using reviewable. Sorry if this is a bit incomprehensible!

@gonzalezjo reviewed 4 files and made 20 comments.
Reviewable status: 4 of 8 files reviewed, 24 unresolved discussions (waiting on hhuangMITRE and jrobble).

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 50 at r3 (raw file):

log = logging.getLogger(__name__)

_LAST_WS_RE = re.compile(r"\s(?=\S*$)")

I would move this to where it's used in _divide() and give it a name like _LAST_WHITESPACE_REGEX. Or maybe even, since we only use it once, just inline its use and leave an explanatory comment stating what the regex does. It just doesn't need to be at the top-level scope.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 159 at r3 (raw file):

            self.sat_model = SaT(sat_model_name)

        # Move model to device; SaT benefits from half precision on GPU.

I assume you mean SaT speed benefits from reduced precision. But reducing precision is a huge behavior change (potentially, at least) w/o QAT quantization-aware training...do we have any data on output quality?

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 210 at r3 (raw file):

        split_mode: str = 'DEFAULT',
        newline_behavior: NewLineBehaviorType = 'GUESS',
        preferred_limit: int = -1

This appears to be entirely undocumented.

(edit: The line I selected in reviewable, specifically, was preferred_limit: int = -1)

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 270 at r3 (raw file):

        substring_list = self._sentence_model.split(substring, lang=self._in_lang)
        if not substring_list:
            return text

Is there any circumstance in which this might happen? (And if so, can we test it?)

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 321 at r3 (raw file):

            else:
                # Split oversized sentence using the default internal logic.
                yield from self._split_sentence_text(sentence)

RE: my next comment, we never enter this branch in tests.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 323 at r3 (raw file):

                yield from self._split_sentence_text(sentence)

    def _split_sentence_text(self, text: str):

This never actually gets run during tests.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 350 at r3 (raw file):

    def _compute_breakpoints_from_sentences(self, text: str, pieces: List[str]) -> List[int]:

This function never gets hit (see the next comment) in tests.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 388 at r3 (raw file):

    def _divide(self, text) -> Tuple[str, str]:
        max_limit = self._limit

I think bad things can happen if limit = 0. Which probably shouldn't ever happen, but I don't know if we maybe want to assert that anywhere, since it would be very hard to debug if what I think could happen (OpenMPF hanging), would happen.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 410 at r3 (raw file):

                    else:
                        left = left_window
                else:

We don't test anything in this branch.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 446 at r3 (raw file):

                cut = len(left)
                if 0 < cut < len(text) and text[cut - 1].isalnum() and text[cut].isalnum():

We don't test this either. My understanding as that this is: "if we're splitting in the middle of a word-like or number-like thing:." So, three things: 1) we should probably try to get coverage of this, (2) we should probably clarify expected behavior somewhere: to me, it's not obvious that we'd treat "999" and "1,000" (or 1.0 and 1; or 1.000 and 1000; etc) in very different ways, but "," would be a valid split here. That's probably fine and there's probably not much we can do to make splitting perfect w/rt commas, but since it's not obvious / may be a bit surprising that that can happen, maybe it can be noted?

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 447 at r3 (raw file):

                cut = len(left)
                if 0 < cut < len(text) and text[cut - 1].isalnum() and text[cut].isalnum():
                    m = _LAST_WS_RE.search(left)

See earlier comments about _LAST_WS_RE and coverage of the inside of this branch.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 449 at r3 (raw file):

                    m = _LAST_WS_RE.search(left)
                    if m:
                        left = left[:m.end()]

Again, it'd be good to get coverage of this case (and the implicit no-op else case)

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 453 at r3 (raw file):

                # Worst-case, but extremely unlikely to happen.
                if left == "" and text != "":
                    left = text[:1]

Again, coverage.

---

detection/nlp_text_splitter/nlp_text_splitter/__init__.py line 457 at r3 (raw file):

                return left, text[len(left):]

            char_per_size = len(left_window) / max(left_size, 1)

Nothing after this line gets coverage (in this function, at least)

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 102 at r3 (raw file):

        # Default to GUESS if None or invalid string
        if behavior is None:
            behavior = 'GUESS'

We don't get any coverage for this in tests - which is fine, but the behavior should probably be in a docstring or we should be using Python's default arguments or something so that at least it's clear that there's a spec we're adhering to w/this function.

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 109 at r3 (raw file):

            return lambda s, l: cls._replace_new_lines(s, cls._guess_lang_separator(s, l))
        elif behavior == 'REMOVE':
            return lambda s, _: cls._replace_new_lines(s, '')

We don't get any coverage for this in tests.

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 111 at r3 (raw file):

            return lambda s, _: cls._replace_new_lines(s, '')
        elif behavior == 'SPACE':
            return lambda s, _: cls._replace_new_lines(s, ' ')

We don't get any coverage for this in tests.

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 115 at r3 (raw file):

            return lambda s, _: s
        else:
            raise mpf.DetectionError.INVALID_PROPERTY.exception(

We don't get any coverage for this in tests.

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 128 at r3 (raw file):

        else:
            first_alpha_letter = next((ch for ch in text if ch.isalpha()), 'a')
            if ChineseAndJapaneseCodePoints.check_char(first_alpha_letter):

What's the premise behind this in a world where we seemingly know the language already and these languages are both covered by NO_SPACE_LANGS? As a related issue: we don't get any coverage of "if language.upper() in NO_SPACE_LANGS:" .

---

detection/nlp_text_splitter/nlp_text_splitter/newline_behavior.py line 150 at r3 (raw file):

            if match_text == '\n':
                # Surrounding characters are not whitespace.
                return replacement

We don't get any coverage for this in tests.