From 89685395703265225d88e6c32eb3b4a18e9081ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Dentan?= <65963192+DentanJeremie@users.noreply.github.com> Date: Sun, 14 Jun 2026 08:40:16 +0200 Subject: [PATCH 1/3] Handled two edge cases: when there is no nltk-tokenizable word in the generation and when there is not EOS and the last char is a punctuation mark. --- pyproject.toml | 2 +- src/much_segmenter/much_segmenter.py | 27 ++++++++--- tests/data.py | 66 ++++++++++++++++++++++++++- tests/test_segmentation.py | 67 +++++++++++++++++++++++++++- 4 files changed, 152 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 14f6b8c..1949850 100644 --- a/pyproject.toml +++ b/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"}, diff --git a/src/much_segmenter/much_segmenter.py b/src/much_segmenter/much_segmenter.py index 2929e52..4cd91d5 100644 --- a/src/much_segmenter/much_segmenter.py +++ b/src/much_segmenter/much_segmenter.py @@ -53,7 +53,7 @@ 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 = [] @@ -61,11 +61,12 @@ def much_segmentation( 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 @@ -73,7 +74,8 @@ def much_segmentation( 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 diff --git a/tests/data.py b/tests/data.py index 19b94e0..5546eca 100644 --- a/tests/data.py +++ b/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", "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", "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", + [[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", + [[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", + [[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", + [[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", + [[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", + [[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", + [[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", + [[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", + [[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", + [[0, 1, 2, 3, 4], [5, 6], [7, 8], [9, 10, 11, 12, 13, 14, 15, 16]]), +] non_idempotent_segmentations = [ ( diff --git a/tests/test_segmentation.py b/tests/test_segmentation.py index ac8dc43..711e541 100644 --- a/tests/test_segmentation.py +++ b/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}" + ) From 9221e0418f889f2ff6167dcf71db82122b53cbb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Dentan?= <65963192+DentanJeremie@users.noreply.github.com> Date: Sun, 14 Jun 2026 08:45:24 +0200 Subject: [PATCH 2/3] Changelog --- README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/README.md b/README.md index b1152cc..822bb43 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,13 @@ 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. From 6dd6885a1be962675c765aada44d3ea58fc50ea3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9r=C3=A9mie=20Dentan?= <65963192+DentanJeremie@users.noreply.github.com> Date: Sun, 14 Jun 2026 08:46:33 +0200 Subject: [PATCH 3/3] Readme --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 822bb43..c6bb5db 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,8 @@ This work received financial support from the research chair *Trustworthy and Re 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.