Fix: No tokenizable word + merge last punctuation into last claim

README.md

@@ -74,12 +74,21 @@ This package is released alongside the MUCH benchmark. This benchmark includes t
   - Generation configs: [orailix/MUCH-configs](https://huggingface.co/datasets/orailix/MUCH-configs)
   - Baseline evaluation data: [orailix/MUCH-signals](https://huggingface.co/datasets/orailix/MUCH-signals)
 
+## Changelog
+
+- v0.3.0 Handles two edge case:
+  - Avoids raising an error when there are no nltk-tokenizable content in the input.
+  - When there are no EOS token and the last char is a punctuation mark, it is merged into the latest claim.
+- v0.2.1 First public version. This version was used to generate MUCH [arXiv:2511.17081](https://arxiv.org/abs/2511.17081)
+
 ## Acknowledgement
 
 This work received financial support from the research chair *Trustworthy and Responsible AI* at École Polytechnique.
 
 This work was granted access to the HPC resources of IDRIS under the allocation **AD011014843R1**, made by GENCI.
 
+We thank Mahammed El Sharkawy for his contribution to this work.
+
 ## Copyright and License
 
 Copyright 2025–present Laboratoire d’Informatique de l’École Polytechnique.

pyproject.toml

@@ -8,7 +8,7 @@ package-dir = {"" = "src"}
 
 [project]
 name = "much-segmenter"
-version = "0.2.1"
+version = "0.3.0"
 license = { text = "Apache-2.0" }
 authors = [
     { name="Jérémie Dentan"},

src/much_segmenter/much_segmenter.py

@@ -53,27 +53,29 @@ def much_segmentation(
     segmentation_punct_set = get_punctuation_set()
 
     # NLTK tokenization + stopwords detection
-    word_indices = nltk_tokenizer.span_tokenize(generation)
+    word_indices = list(nltk_tokenizer.span_tokenize(generation))
 
     # Init containers for loop
     char_idx_of_claim_start = []
     previous_word_was_stopword = False
     next_word_is_stopword = False
 
     # Index of EOS
+    eos_index = None
     for final_token in get_known_eos_or_eot(llm_tokenizer):
         if llm_tokenizer.decode(final_token) in generation:
             eos_index = generation.index(llm_tokenizer.decode(final_token))
             break
-    else:
+    if eos_index is None:
         eos_index = len(generation)
 
     # Iterating over NLTK words - the goal here is to fill char_idx_of_claim_start
     # We ommit "0" in char_idx_of_claim_start on purpose, to integrate with the LLM token part
     for word_start, word_stop in word_indices:
         # Did we reach EOS ?
         if word_stop >= eos_index:
-            char_idx_of_claim_start.append(eos_index)
+            if eos_index < len(generation):
+                char_idx_of_claim_start.append(eos_index)
             break
 
         # Is current word a stopword?
@@ -83,7 +85,6 @@ def much_segmentation(
             or (current_word in segmentation_punct_set)
             or next_word_is_stopword
         )
-
         # We only add it if the previous word was not a stopword
         if word_is_stopword and (not previous_word_was_stopword) and (word_start != 0):
             char_idx_of_claim_start.append(word_start)
@@ -93,8 +94,8 @@ def much_segmentation(
         previous_word_was_stopword = word_is_stopword
         next_word_is_stopword = current_word[-1] == "."
 
-    # Adding the latest part to terminate chunks
-    if char_idx_of_claim_start[-1] != len(generation):
+    # Adding the last part to terminate chunks.
+    if not char_idx_of_claim_start or char_idx_of_claim_start[-1] != len(generation):
         char_idx_of_claim_start.append(len(generation))
 
     # Now, we iterate on the LLM tokens and detect when we "cross" a char in char_idx_of_claim_start
@@ -125,4 +126,18 @@ def much_segmentation(
     if current_chunk != []:
         result.append(current_chunk)
 
+    # Post-processing
+    # If the last token is punctuation, include it in the final chunk.
+    if len(result) >= 2:
+        last_is_eos = len(result[-1]) == 1 and output_tokens[
+            result[-1][0]
+        ] in get_known_eos_or_eot(llm_tokenizer)
+        if not last_is_eos:
+            chunk_text = llm_tokenizer.decode(
+                [output_tokens[i] for i in result[-1]]
+            ).strip()
+            if chunk_text in segmentation_punct_set:
+                result[-2] = result[-2] + result[-1]
+                result.pop()
+
     return result

tests/data.py

@@ -136,8 +136,70 @@
 		[[0, 1, 2], [3, 4, 5, 6], [7, 8, 9, 10], [11, 12, 13, 14, 15, 16, 17], [18, 19, 20, 21, 22, 23, 24]],
 		[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11, 12]],
 		[[0, 1, 2], [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], [14, 15, 16, 17, 18, 19, 20]],
-		[[0, 1, 2], [3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15], [16, 17, 18], [19, 20, 21, 22], [23, 24, 25], [26, 27], [28, 29, 30], [31, 32, 33, 34, 35, 36, 37]],    ]
-    }
+		[[0, 1, 2], [3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15], [16, 17, 18], [19, 20, 21, 22], [23, 24, 25], [26, 27], [28, 29, 30], [31, 32, 33, 34, 35, 36, 37]],
+    ]
+}
+
+# the EOS token in the phrase is recognised by that model :
+# the punctuation chunk immediately before the EOS must remain isolated.
+PUNCT_BEFORE_EOS_PHRASES = [
+    ("The Eiffel Tower was built in 1889?\n</s>",          "mistralai/Ministral-8B-Instruct-2410"),
+    ("The Eiffel Tower was built in 1889?\n<|eot_id|>",    "meta-llama/Llama-3.2-3B-Instruct"),
+    ("The Eiffel Tower was built in 1889?\n<|eot_id|>",    "meta-llama/Llama-3.1-8B-Instruct"),
+    ("The Eiffel Tower was built in 1889?\n<end_of_turn>", "google/gemma-3-4b-it"),
+]
+
+
+#
+# Three categories:
+#   No EOS token, short phrase : eos_index == len(generation), single chunk
+#   No EOS token + space : post-processing merges isolated terminal "."
+#   Punct immediately before a recognised EOS : last_is_eos guard keeps it isolated
+edge_case_segmentations = [
+    # Category 1
+    ("mistralai/Ministral-8B-Instruct-2410", "Unknown.", [[0, 1]]),
+    ("meta-llama/Llama-3.2-3B-Instruct",     "Unknown.", [[0, 1]]),
+    ("meta-llama/Llama-3.1-8B-Instruct",     "Unknown.", [[0, 1]]),
+    ("google/gemma-3-4b-it",                  "Unknown.", [[0, 1]]),
+    ("bert-base-uncased",                     "Unknown.", [[0, 1]]),
+    # Category 2
+    ("mistralai/Ministral-8B-Instruct-2410", "Unknown. ", [[0, 1, 2]]),
+    ("meta-llama/Llama-3.2-3B-Instruct",     "Unknown. ", [[0, 1, 2]]),
+    ("meta-llama/Llama-3.1-8B-Instruct",     "Unknown. ", [[0, 1, 2]]),
+    ("google/gemma-3-4b-it",                  "Unknown. ", [[0, 1, 2]]),
+    ("bert-base-uncased",                     "Unknown. ", [[0, 1]]),
+    # Category 3
+    ("mistralai/Ministral-8B-Instruct-2410", "The Eiffel Tower was built in 1889?\n</s>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10, 11, 12], [13], [14]]),
+    ("meta-llama/Llama-3.2-3B-Instruct",     "The Eiffel Tower was built in 1889?\n</s>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12, 13, 14]]),
+    ("meta-llama/Llama-3.1-8B-Instruct",     "The Eiffel Tower was built in 1889?\n</s>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12, 13, 14]]),
+    ("google/gemma-3-4b-it",                  "The Eiffel Tower was built in 1889?\n</s>",
+        [[0, 1, 2], [3, 4], [5, 6, 7, 8, 9, 10], [11, 12, 13]]),
+    ("bert-base-uncased",                     "The Eiffel Tower was built in 1889?\n</s>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8], [9, 10, 11, 12, 13]]),
+    ("mistralai/Ministral-8B-Instruct-2410", "The Eiffel Tower was built in 1889?\n<|eot_id|>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10, 11, 12], [13, 14, 15, 16, 17, 18, 19, 20]]),
+    ("meta-llama/Llama-3.2-3B-Instruct",     "The Eiffel Tower was built in 1889?\n<|eot_id|>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11], [12]]),
+    ("meta-llama/Llama-3.1-8B-Instruct",     "The Eiffel Tower was built in 1889?\n<|eot_id|>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11], [12]]),
+    ("google/gemma-3-4b-it",                  "The Eiffel Tower was built in 1889?\n<|eot_id|>",
+        [[0, 1, 2], [3, 4], [5, 6, 7, 8, 9, 10], [11, 12, 13, 14, 15, 16, 17, 18, 19]]),
+    ("bert-base-uncased",                     "The Eiffel Tower was built in 1889?\n<|eot_id|>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8], [9, 10, 11, 12, 13, 14, 15, 16, 17]]),
+    ("mistralai/Ministral-8B-Instruct-2410", "The Eiffel Tower was built in 1889?\n<end_of_turn>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10, 11, 12], [13, 14, 15, 16, 17, 18, 19]]),
+    ("meta-llama/Llama-3.2-3B-Instruct",     "The Eiffel Tower was built in 1889?\n<end_of_turn>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12, 13, 14, 15, 16]]),
+    ("meta-llama/Llama-3.1-8B-Instruct",     "The Eiffel Tower was built in 1889?\n<end_of_turn>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12, 13, 14, 15, 16]]),
+    ("google/gemma-3-4b-it",                  "The Eiffel Tower was built in 1889?\n<end_of_turn>",
+        [[0, 1, 2], [3, 4], [5, 6, 7, 8, 9, 10], [11, 12], [13]]),
+    ("bert-base-uncased",                     "The Eiffel Tower was built in 1889?\n<end_of_turn>",
+        [[0, 1, 2, 3, 4], [5, 6], [7, 8], [9, 10, 11, 12, 13, 14, 15, 16]]),
+]
 
 non_idempotent_segmentations = [
     (

tests/test_segmentation.py

@@ -4,10 +4,17 @@
 from src.much_segmenter import much_segmentation
 from src.much_segmenter.utils import (
     get_known_eos_or_eot,
+    get_punctuation_set,
     get_tokenizer,
     get_tokens_and_spans,
 )
-from tests.data import expected_outputs, non_idempotent_segmentations, test_phrases
+from tests.data import (
+    PUNCT_BEFORE_EOS_PHRASES,
+    edge_case_segmentations,
+    expected_outputs,
+    non_idempotent_segmentations,
+    test_phrases,
+)
 
 
 @pytest.mark.parametrize("model_name", list(expected_outputs.keys()))
@@ -112,3 +119,61 @@ def test_eos_eot_in_its_own_claim(model_name):
         assert len(segmentation[-1]) == 1
 
     assert model_name == "bert-base-uncased" or count_test > 0
+
+
+@pytest.mark.parametrize("model_name", list(expected_outputs.keys()))
+def test_no_terminal_punct_isolation(model_name):
+
+    tokenizer = get_tokenizer(model_name)
+    punct_set = get_punctuation_set()
+    tokens_all = tokenizer.encode
+
+    for phrase in test_phrases:
+        segmentation = much_segmentation(phrase, tokenizer)
+        encoded = tokenizer.encode(phrase, add_special_tokens=False)
+
+        last_is_eos = len(segmentation[-1]) == 1 and encoded[
+            segmentation[-1][0]
+        ] in get_known_eos_or_eot(tokenizer)
+        if last_is_eos:
+            continue
+
+        chunk_tokens = [encoded[i] for i in segmentation[-1]]
+        chunk_text = tokenizer.decode(chunk_tokens).strip()
+        assert chunk_text not in punct_set, (
+            f"Terminal punctuation isolated in its own chunk for model "
+            f"'{model_name}' on phrase '{phrase}': last content chunk decodes to '{chunk_text}'"
+        )
+
+
+@pytest.mark.parametrize("model_name,phrase,expected", edge_case_segmentations)
+def test_edge_case_segmentation(model_name, phrase, expected):
+    tokenizer = get_tokenizer(model_name)
+    out = much_segmentation(phrase, tokenizer)
+    assert out == expected, (
+        f"Edge-case mismatch for model '{model_name}' on phrase {phrase!r}: "
+        f"got {out}, expected {expected}"
+    )
+
+
+@pytest.mark.parametrize("phrase,model_name", PUNCT_BEFORE_EOS_PHRASES)
+def test_punct_before_eos_not_merged(phrase, model_name):
+    from src.much_segmenter.utils import get_punctuation_set
+
+    tokenizer = get_tokenizer(model_name)
+    encoded = tokenizer.encode(phrase, add_special_tokens=False)
+    seg = much_segmentation(phrase, tokenizer)
+
+    assert (
+        len(seg) >= 3
+    ), f"Expected at least 3 chunks for '{model_name}' on {phrase!r}, got {seg}"
+
+    assert len(seg[-1]) == 1 and encoded[seg[-1][0]] in get_known_eos_or_eot(
+        tokenizer
+    ), f"Last chunk is not a lone EOS token for '{model_name}' on {phrase!r}"
+
+    pre_eos_text = tokenizer.decode([encoded[i] for i in seg[-2]]).strip()
+    assert pre_eos_text in get_punctuation_set(), (
+        f"Expected punctuation chunk before EOS for '{model_name}' on {phrase!r}, "
+        f"got {pre_eos_text!r}"
+    )