use std::collections::HashSet;
use regex::Regex;
use unicode_segmentation::UnicodeSegmentation;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
use crate::common::{get_special_char_regex, process_word, PUNCTUATION};
pub struct DocumentProcessor<'a> {
documents: &'a [String],
stopwords: HashSet<String>,
punctuation: HashSet<String>,
}
impl<'a> DocumentProcessor<'a> {
pub fn new(
documents: &'a [String],
stopwords: &'a [String],
punctuation: &'a Option<&'a [String]>,
) -> Self {
Self {
documents,
stopwords: stopwords
.iter()
.map(|s| s.to_owned())
.collect::<HashSet<String>>(),
punctuation: punctuation
.unwrap_or(
&PUNCTUATION
.iter()
.map(|s| s.to_string())
.collect::<Vec<String>>(),
)
.iter()
.map(|s| s.to_string())
.collect::<HashSet<String>>(),
}
}
fn process_document(&self, document: &str, special_char_regex: &Option<Regex>) -> String {
document
.unicode_sentences()
.map(|s| {
s.split_word_bounds()
.filter_map(|w| {
process_word(w, special_char_regex, &self.stopwords, &self.punctuation)
})
.collect::<Vec<String>>()
.join(" ")
})
.collect::<Vec<String>>()
.join(" ")
}
pub fn process_documents(&self) -> Vec<String> {
let special_char_regex = get_special_char_regex();
#[cfg(feature = "parallel")]
{
self.documents
.par_iter()
.map(|doc| self.process_document(doc, &special_char_regex))
.collect::<Vec<String>>()
}
#[cfg(not(feature = "parallel"))]
{
self.documents
.iter()
.map(|doc| self.process_document(doc, &special_char_regex))
.collect::<Vec<String>>()
}
}
}