unobtanium_segmenter/segmentation/

lingua_with_sentence_splitter.rs

// SPDX-FileCopyrightText: 2026 Slatian
//
// SPDX-License-Identifier: LGPL-3.0-only

use std::vec::IntoIter;

use lingua::LanguageDetector;
use lingua::LanguageDetectorBuilder;
use unicode_segmentation::UnicodeSegmentation;

use crate::SegmentedToken;
use crate::UseOrSubdivide;
use crate::language_detection::DetectionIterator;
use crate::language_detection::detect_top_n_languages;
use crate::language_detection::lingua_language_to_whatlang_language;
use crate::segmentation::Segmenter;
use crate::segmentation::UnicodeSentenceSplitter;

/// Language splitter and detector.
///
/// This is intended to be used near/at the start of the segmentation chain.
///
/// Performance: There is a significant spin-up cost involved when using this struct for the first time that won't happen on subsequent uses. This only happens once per program, independent of you keeping a specific instance or not. See the [lingua] documentation for RAM requirements.
///
/// Using low accuracy mode probably isn't worth it.
///
/// Language support: Currently this crate uses [whatlang::Lang] to communicate languages, so the language support is the intersection of whatlang and what [lingua::Language] both support.
pub struct LinguaLanguageBlockSentenceSplitter {
	language_detector: LanguageDetector,
}

impl Default for LinguaLanguageBlockSentenceSplitter {
	fn default() -> Self {
		Self {
			language_detector: LanguageDetectorBuilder::from_all_languages()
				.with_preloaded_language_models()
				.build(),
		}
	}
}

impl LinguaLanguageBlockSentenceSplitter {
	/// Create a new LiguaLanguageBlockSplitter instance that is configured to preload all languages on the first use.
	pub fn new() -> Self {
		Default::default()
	}

	/// Create a new LiguaLanguageBlockSplitter from a custom [LanguageDetectorBuilder].
	pub fn new_with_builder(mut builder: LanguageDetectorBuilder) -> Self {
		Self {
			language_detector: builder.build(),
		}
	}
}

impl Segmenter for LinguaLanguageBlockSentenceSplitter {
	type SubdivisionIter<'a> = IntoIter<SegmentedToken<'a>>;

	fn subdivide<'a>(
		&self,
		token: SegmentedToken<'a>,
	) -> UseOrSubdivide<SegmentedToken<'a>, IntoIter<SegmentedToken<'a>>> {
		let languages = detect_top_n_languages(&self.language_detector, 3, token.text);

		println!("Languages: {languages:?}");

		if languages.is_empty() {
			// No lnaguage found: Only do sentence splitting
			return UnicodeSentenceSplitter::new().subdivide(token);
		}

		let detection_iterator = DetectionIterator::detect(
			// This will panic for any empty languages vec, keep the check above intact
			&LanguageDetectorBuilder::from_languages(&languages).build(),
			token.text,
		);
		let mut sentence_iterator = token.text.split_sentence_bounds().map(|s| s.len());

		let mut collection: Vec<SegmentedToken<'_>> = Vec::new();

		let mut last_offset = 0;
		let mut next_sentence_split_at = sentence_iterator.next().unwrap_or(token.text.len());

		for (mut start_index, end_index, language) in detection_iterator {
			println!(
				"{start_index}..{end_index} next split: {next_sentence_split_at}, lang: {language:?}"
			);
			let language = language.and_then(lingua_language_to_whatlang_language);
			if next_sentence_split_at > start_index && next_sentence_split_at < end_index {
				collection.push(
					SegmentedToken::new_derived_from(
						&token.text[start_index..next_sentence_split_at],
						&token,
					)
					.with_detected_language(language, true, 1.0),
				);
				collection.push(
					SegmentedToken::new_end_of_sentence(
						&token.text[next_sentence_split_at..next_sentence_split_at],
					)
					.with_detected_language(language, true, 1.0),
				);

				start_index = next_sentence_split_at;
				next_sentence_split_at = sentence_iterator
					.next()
					.map(|n| n + next_sentence_split_at)
					.unwrap_or_else(|| token.text.len());
			}
			collection.push(
				SegmentedToken::new_derived_from(&token.text[start_index..end_index], &token)
					.with_detected_language(language, true, 1.0),
			);

			if next_sentence_split_at == end_index {
				collection.push(
					SegmentedToken::new_end_of_sentence(
						&token.text[next_sentence_split_at..next_sentence_split_at],
					)
					.with_detected_language(language, true, 1.0),
				);
				next_sentence_split_at = sentence_iterator
					.next()
					.map(|n| n + next_sentence_split_at)
					.unwrap_or_else(|| token.text.len());
			}

			last_offset = end_index;
		}

		if last_offset != token.text.len() {
			collection.push(SegmentedToken::new_derived_from(
				&token.text[last_offset..],
				&token,
			));
		}

		UseOrSubdivide::Subdivide(collection.into_iter())
	}
}

#[cfg(test)]
mod test {

	use whatlang::Lang;

	use std::time::Instant;

	use super::*;

	use crate::chain::ChainSegmenter;
	use crate::chain::StartSegmentationChain;

	#[test]
	fn test_lingua_multilanguage_detection() {
		let test_text = "Parlez-vous français? \
			Ich spreche Französisch nur ein bisschen. \
			A little bit is better than nothing.   ";

		let lingua_segmenter = LinguaLanguageBlockSentenceSplitter::new();

		for _ in 0..100 {
			let result: Vec<(&str, Option<Lang>)> = test_text
				.start_segmentation_chain()
				.chain_segmenter(&lingua_segmenter)
				.map(|t| (t.text, t.detected_language))
				.collect();

			let expected_tokens = vec![
				("Parlez-vous français? ", Some(Lang::Fra)),
				("", Some(Lang::Fra)),
				(
					"Ich spreche Französisch nur ein bisschen. ",
					Some(Lang::Deu),
				),
				("", Some(Lang::Deu)),
				("A little bit is better than nothing.   ", Some(Lang::Eng)),
				("", Some(Lang::Eng)),
			];

			assert_eq!(result, expected_tokens);
		}
	}

	#[test]
	fn test_lingua_performance() {
		let test_text = "Parlez-vous français? \
    Ich spreche Französisch nur ein bisschen. \
    A little bit is better than nothing.";

		let start_instant = Instant::now();

		let _result: Vec<(&str, Option<Lang>)> = test_text
			.start_segmentation_chain()
			.chain_segmenter(&LinguaLanguageBlockSentenceSplitter::new())
			.map(|t| (t.text, t.detected_language))
			.collect();

		let time_first_iteration = start_instant.elapsed();
		let start_instant = Instant::now();

		for _ in 0..100 {
			let _result: Vec<(&str, Option<Lang>)> = test_text
				.start_segmentation_chain()
				.chain_segmenter(&LinguaLanguageBlockSentenceSplitter::new())
				.map(|t| (t.text, t.detected_language))
				.collect();
		}

		let time_multiple_iterations = start_instant.elapsed();

		assert!(
			time_first_iteration > (time_multiple_iterations / 100),
			"Subsequent iterations should be faster than the initial one, even when not keeping the struct around. {time_first_iteration:?} > {:?}",
			time_multiple_iterations / 100
		);
	}
}