Using Rust-Bert to Summarize Blog Posts

Create a BART Summarizer in rust-bert

rust-bert makes using pre-trained models very easy from Rust.

use rust_bert::bart::{BartConfigResources, BartMergesResources, BartModelResources, BartVocabResources};
use rust_bert::pipelines::summarization::{SummarizationConfig, SummarizationModel};
use rust_bert::resources::{RemoteResource, Resource};
use rust_bert::{RustBertError};
use tch::{Device};
use std::sync::Mutex;
use lazy_static::lazy_static;

lazy_static! {
   static ref summarization_model: Mutex<SummarizationModel> = {
      let config_resource = Resource::Remote(RemoteResource::from_pretrained(
         BartConfigResources::DISTILBART_CNN_6_6,
      ));
      let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
         BartVocabResources::DISTILBART_CNN_6_6,
      ));
      let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
         BartMergesResources::DISTILBART_CNN_6_6,
      ));
      let model_resource = Resource::Remote(RemoteResource::from_pretrained(
         BartModelResources::DISTILBART_CNN_6_6,
      ));
      let summarization_config = SummarizationConfig {
         model_resource,
         config_resource,
         vocab_resource,
         merges_resource,
         num_beams: 1,
         length_penalty: 2.0,
         device: Device::Cpu,
         ..Default::default()
      };

      let model = SummarizationModel::new(summarization_config).expect("Error loading summarization model");
      Mutex::new(model)
   };
}

pub fn summarize(input: &str) -> String {
   impl_summarize(input).unwrap_or("".to_string())
}
fn impl_summarize(input: &str) -> Result<String,RustBertError> {
   let output = summarization_model.lock().unwrap().summarize(&[input]);
   Ok(output[0].to_string())
}