Struct KeywordExtractionModel

pub struct KeywordExtractionModel<'a> {
    pub sentence_embeddings_model: SentenceEmbeddingsModel,
    pub tokenizer: StopWordsTokenizer<'a>,
    /* private fields */
}

KeywordExtractionModel to extract keywords from input texts

It contains a sentence embeddings model to compute word-document similarities, a tokenizer to define a keyword candidates list and a scorer to rank these keywords.

• sentence_embeddings_model: Sentence embeddings model
• tokenizer: tokenizer used to generate the list of candidates (differs from the transformer tokenizer)

Fields

sentence_embeddings_model: SentenceEmbeddingsModel

tokenizer: StopWordsTokenizer<'a>

Implementations

impl<'a> KeywordExtractionModel<'a>

pub fn new( config: KeywordExtractionConfig<'a>, ) -> Result<KeywordExtractionModel<'a>, RustBertError>

Build a new KeywordExtractionModel

Arguments

• config - KeywordExtractionConfig object containing a sentence embeddings configuration and tokenizer-specific options

Example

use rust_bert::pipelines::keywords_extraction::KeywordExtractionModel;

let keyword_extraction_model = KeywordExtractionModel::new(Default::default())?;

pub fn get_tokenizer(&self) -> &TokenizerOption

Get a reference to the model tokenizer.

pub fn get_tokenizer_mut(&mut self) -> &mut TokenizerOption

Get a mutable reference to the model tokenizer.

pub fn predict<S>( &self, inputs: &[S], ) -> Result<Vec<Vec<Keyword>>, RustBertError>

where S: AsRef<str> + Send + Sync,

Extract keywords from a list of input texts.

Arguments

• inputs - slice of string-like input texts to extract keywords from

Returns

• Result<Vec<Vec<Keyword>>, RustBertError> containing a list of keyword for each input text

Example

use rust_bert::pipelines::keywords_extraction::KeywordExtractionModel;

let keyword_extraction_model = KeywordExtractionModel::new(Default::default())?;
let input = [
    "This is a first sentence to extract keywords from.",
    "Some keywords will be extracted from this text too.",
];
let output = keyword_extraction_model.predict(&input);