use std::collections::HashMap;
mod text_rank_logic;
pub mod text_rank_params;
use text_rank_logic::TextRankLogic;
pub use text_rank_params::TextRankParams;
use crate::{
common::{get_ranked_scores, get_ranked_strings},
tokenizer::Tokenizer,
};
pub struct TextRank {
word_rank: HashMap<String, f32>,
phrase_rank: HashMap<String, f32>,
}
impl TextRank {
pub fn new(params: TextRankParams) -> Self {
let (text, stop_words, punctuation, window_size, damping, tol, phrase_length) =
params.get_params();
let tokenizer = Tokenizer::new(text, stop_words, punctuation);
let (word_rank, phrase_rank) = TextRankLogic::build_text_rank(
tokenizer.sync_split_into_words(),
tokenizer.sync_split_into_phrases(phrase_length),
window_size,
damping,
tol,
);
Self {
word_rank,
phrase_rank,
}
}
pub fn get_word_score(&self, word: &str) -> f32 {
*self.word_rank.get(word).unwrap_or(&0.0)
}
pub fn get_phrase_score(&self, phrase: &str) -> f32 {
*self.phrase_rank.get(phrase).unwrap_or(&0.0)
}
pub fn get_ranked_words(&self, n: usize) -> Vec<String> {
get_ranked_strings(&self.word_rank, n)
}
pub fn get_ranked_word_scores(&self, n: usize) -> Vec<(String, f32)> {
get_ranked_scores(&self.word_rank, n)
}
pub fn get_ranked_phrases(&self, n: usize) -> Vec<String> {
get_ranked_strings(&self.phrase_rank, n)
}
pub fn get_ranked_phrase_scores(&self, n: usize) -> Vec<(String, f32)> {
get_ranked_scores(&self.phrase_rank, n)
}
pub fn get_word_scores_map(&self) -> &HashMap<String, f32> {
&self.word_rank
}
pub fn get_phrase_scores_map(&self) -> &HashMap<String, f32> {
&self.phrase_rank
}
}