mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
318 lines
14 KiB
Rust
318 lines
14 KiB
Rust
extern crate failure;
|
|
|
|
use rust_bert::albert::{AlbertConfigResources, AlbertModelResources, AlbertVocabResources};
|
|
use rust_bert::bart::{
|
|
BartConfigResources, BartMergesResources, BartModelResources, BartVocabResources,
|
|
};
|
|
use rust_bert::bert::{BertConfigResources, BertModelResources, BertVocabResources};
|
|
use rust_bert::distilbert::{
|
|
DistilBertConfigResources, DistilBertModelResources, DistilBertVocabResources,
|
|
};
|
|
use rust_bert::electra::{ElectraConfigResources, ElectraModelResources, ElectraVocabResources};
|
|
use rust_bert::gpt2::{
|
|
Gpt2ConfigResources, Gpt2MergesResources, Gpt2ModelResources, Gpt2VocabResources,
|
|
};
|
|
use rust_bert::openai_gpt::{
|
|
OpenAiGptConfigResources, OpenAiGptMergesResources, OpenAiGptModelResources,
|
|
OpenAiGptVocabResources,
|
|
};
|
|
use rust_bert::resources::{download_resource, RemoteResource, Resource};
|
|
use rust_bert::roberta::{
|
|
RobertaConfigResources, RobertaMergesResources, RobertaModelResources, RobertaVocabResources,
|
|
};
|
|
|
|
/// This example downloads and caches all dependencies used in model tests. This allows for safe
|
|
/// multi threaded testing (two test using the same resource would otherwise download the file to
|
|
/// the same location).
|
|
|
|
fn download_distil_gpt2() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the HuggingFace Inc. team at https://huggingface.co/models. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2ConfigResources::DISTIL_GPT2,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2VocabResources::DISTIL_GPT2,
|
|
));
|
|
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2MergesResources::DISTIL_GPT2,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2ModelResources::DISTIL_GPT2,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_distilbert_sst2() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the HuggingFace Inc. team at https://huggingface.co/models. Modified with conversion to C-array format.
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertModelResources::DISTIL_BERT_SST2,
|
|
));
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertConfigResources::DISTIL_BERT_SST2,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertVocabResources::DISTIL_BERT_SST2,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_distilbert_qa() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the HuggingFace Inc. team at https://huggingface.co/models. Modified with conversion to C-array format.
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertModelResources::DISTIL_BERT_SQUAD,
|
|
));
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertConfigResources::DISTIL_BERT_SQUAD,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertVocabResources::DISTIL_BERT_SQUAD,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_distilbert() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the HuggingFace Inc. team at https://huggingface.co/models. Modified with conversion to C-array format.
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertModelResources::DISTIL_BERT,
|
|
));
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertConfigResources::DISTIL_BERT,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
DistilBertVocabResources::DISTIL_BERT,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_gpt2() -> failure::Fallible<()> {
|
|
// Shared under Modified MIT license by the OpenAI team at https://github.com/openai/gpt-2. Modified with conversion to C-array format.
|
|
let config_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(Gpt2ConfigResources::GPT2));
|
|
let vocab_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(Gpt2VocabResources::GPT2));
|
|
let merges_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(Gpt2MergesResources::GPT2));
|
|
let weights_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(Gpt2ModelResources::GPT2));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_gpt() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the OpenAI team at https://github.com/openai/finetune-transformer-lm. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
OpenAiGptConfigResources::GPT,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
OpenAiGptVocabResources::GPT,
|
|
));
|
|
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
OpenAiGptMergesResources::GPT,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
OpenAiGptModelResources::GPT,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_roberta() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the Facebook AI Research Fairseq team at https://github.com/pytorch/fairseq. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
RobertaConfigResources::ROBERTA,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
RobertaVocabResources::ROBERTA,
|
|
));
|
|
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
RobertaMergesResources::ROBERTA,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
RobertaModelResources::ROBERTA,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_bert() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the Google team at https://github.com/google-research/bert. Modified with conversion to C-array format.
|
|
let config_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BertConfigResources::BERT));
|
|
let vocab_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BertVocabResources::BERT));
|
|
let weights_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BertModelResources::BERT));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_bert_ner() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the MDZ Digital Library team at the Bavarian State Library at https://github.com/dbmdz/berts. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BertConfigResources::BERT_NER,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BertVocabResources::BERT_NER,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BertModelResources::BERT_NER,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_bart() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the Facebook AI Research Fairseq team at https://github.com/pytorch/fairseq. Modified with conversion to C-array format.
|
|
let config_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BartConfigResources::BART));
|
|
let vocab_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BartVocabResources::BART));
|
|
let merges_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BartMergesResources::BART));
|
|
let weights_resource =
|
|
Resource::Remote(RemoteResource::from_pretrained(BartModelResources::BART));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_bart_cnn() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the Facebook AI Research Fairseq team at https://github.com/pytorch/fairseq. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BartConfigResources::BART_CNN,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BartVocabResources::BART_CNN,
|
|
));
|
|
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BartMergesResources::BART_CNN,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
BartModelResources::BART_CNN,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_electra_generator() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the Google team at https://github.com/google-research/electra. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraConfigResources::BASE_GENERATOR,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraVocabResources::BASE_GENERATOR,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraModelResources::BASE_GENERATOR,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_electra_discriminator() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the Google team at https://github.com/google-research/electra. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraConfigResources::BASE_DISCRIMINATOR,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraVocabResources::BASE_DISCRIMINATOR,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
ElectraModelResources::BASE_DISCRIMINATOR,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn download_albert_base_v2() -> failure::Fallible<()> {
|
|
// Shared under Apache 2.0 license by the Google team at https://github.com/google-research/ALBERT. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
AlbertConfigResources::ALBERT_BASE_V2,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
AlbertVocabResources::ALBERT_BASE_V2,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
AlbertModelResources::ALBERT_BASE_V2,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn _download_dialogpt() -> failure::Fallible<()> {
|
|
// Shared under MIT license by the Microsoft team at https://huggingface.co/microsoft/DialoGPT-medium. Modified with conversion to C-array format.
|
|
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2ConfigResources::DIALOGPT_MEDIUM,
|
|
));
|
|
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2VocabResources::DIALOGPT_MEDIUM,
|
|
));
|
|
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2MergesResources::DIALOGPT_MEDIUM,
|
|
));
|
|
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
|
|
Gpt2ModelResources::DIALOGPT_MEDIUM,
|
|
));
|
|
let _ = download_resource(&config_resource)?;
|
|
let _ = download_resource(&vocab_resource)?;
|
|
let _ = download_resource(&merges_resource)?;
|
|
let _ = download_resource(&weights_resource)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn main() -> failure::Fallible<()> {
|
|
let _ = download_distil_gpt2();
|
|
let _ = download_distilbert_sst2();
|
|
let _ = download_distilbert_qa();
|
|
let _ = download_distilbert();
|
|
let _ = download_gpt2();
|
|
let _ = download_gpt();
|
|
let _ = download_roberta();
|
|
let _ = download_bert();
|
|
let _ = download_bert_ner();
|
|
let _ = download_bart();
|
|
let _ = download_bart_cnn();
|
|
let _ = download_electra_generator();
|
|
let _ = download_electra_discriminator();
|
|
let _ = download_albert_base_v2();
|
|
|
|
Ok(())
|
|
}
|