guillaume-be/rust-bert
1
1
Fork 0
mirror of https://github.com/guillaume-be/rust-bert.git synced 2024-09-11 12:55:34 +03:00
Code Issues Packages Projects Releases Wiki Activity

Updated BART embeddings for compatibility with Marian model

Browse Source
This commit is contained in:
Guillaume B 2021-03-27 18:03:45 +01:00
parent 32537da610
commit 5e6b84a7a0
1 changed files with 1 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
examples/summarization.rs
View File

@ -21,37 +21,7 @@ use rust_bert::resources::{RemoteResource, Resource};
use tch::Device;
fn main() -> anyhow::Result<()> {
// let summarization_model = SummarizationModel::new(Default::default())?;
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
BartConfigResources::DISTILBART_CNN_12_6,
));
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
BartVocabResources::DISTILBART_CNN_12_6,
));
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
BartMergesResources::DISTILBART_CNN_12_6,
));
let model_resource = Resource::Remote(RemoteResource::from_pretrained(
BartModelResources::DISTILBART_CNN_12_6,
));
let summarization_config = SummarizationConfig {
model_type: ModelType::Bart,
model_resource,
config_resource,
vocab_resource,
merges_resource,
length_penalty: 1.0,
num_beams: 4,
no_repeat_ngram_size: 3,
min_length: 56,
max_length: 142,
device: Device::cuda_if_available(),
..Default::default()
};
let summarization_model = SummarizationModel::new(summarization_config)?;
let summarization_model = SummarizationModel::new(Default::default())?;
let input = ["In findings published Tuesday in Cornell University's arXiv by a team of scientists \
from the University of Montreal and a separate report published Wednesday in Nature Astronomy by a team \