mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
CodeBERT Pretrained models and examples (#322)
* Addition of Codebert examples * Addition of CodeBERT pretrained models, CodeBERT example
This commit is contained in:
parent
f12e8ef475
commit
0fc5ce6ad4
@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file. The format
|
|||||||
- Addition of All-MiniLM-L6-V2 model weights
|
- Addition of All-MiniLM-L6-V2 model weights
|
||||||
- Addition of Keyword/Keyphrases extraction pipeline based on KeyBERT (https://github.com/MaartenGr/KeyBERT)
|
- Addition of Keyword/Keyphrases extraction pipeline based on KeyBERT (https://github.com/MaartenGr/KeyBERT)
|
||||||
- Addition of Masked Language Model pipeline, allowing to predict masked words.
|
- Addition of Masked Language Model pipeline, allowing to predict masked words.
|
||||||
|
- Support for the CodeBERT language model with pretrained models for language detection and masked token prediction
|
||||||
|
|
||||||
## Changed
|
## Changed
|
||||||
- Addition of type aliases for the controlled generation (`PrefixAllowedFunction`) and zero-shot classification (`ZeroShotTemplate`).
|
- Addition of type aliases for the controlled generation (`PrefixAllowedFunction`) and zero-shot classification (`ZeroShotTemplate`).
|
||||||
|
85
examples/codebert.rs
Normal file
85
examples/codebert.rs
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
// Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
|
||||||
|
// Copyright 2019 Guillaume Becquin
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
extern crate anyhow;
|
||||||
|
|
||||||
|
use rust_bert::pipelines::common::ModelType;
|
||||||
|
use rust_bert::pipelines::masked_language::{MaskedLanguageConfig, MaskedLanguageModel};
|
||||||
|
use rust_bert::pipelines::sequence_classification::{
|
||||||
|
SequenceClassificationConfig, SequenceClassificationModel,
|
||||||
|
};
|
||||||
|
use rust_bert::resources::RemoteResource;
|
||||||
|
use rust_bert::roberta::{
|
||||||
|
RobertaConfigResources, RobertaMergesResources, RobertaModelResources, RobertaVocabResources,
|
||||||
|
};
|
||||||
|
|
||||||
|
fn main() -> anyhow::Result<()> {
|
||||||
|
// Language identification
|
||||||
|
let sequence_classification_config = SequenceClassificationConfig::new(
|
||||||
|
ModelType::Roberta,
|
||||||
|
RemoteResource::from_pretrained(RobertaModelResources::CODEBERTA_LANGUAGE_ID),
|
||||||
|
RemoteResource::from_pretrained(RobertaConfigResources::CODEBERTA_LANGUAGE_ID),
|
||||||
|
RemoteResource::from_pretrained(RobertaVocabResources::CODEBERTA_LANGUAGE_ID),
|
||||||
|
Some(RemoteResource::from_pretrained(
|
||||||
|
RobertaMergesResources::CODEBERTA_LANGUAGE_ID,
|
||||||
|
)),
|
||||||
|
false,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
let sequence_classification_model =
|
||||||
|
SequenceClassificationModel::new(sequence_classification_config)?;
|
||||||
|
|
||||||
|
// Define input
|
||||||
|
let input = [
|
||||||
|
"def f(x):\
|
||||||
|
return x**2",
|
||||||
|
"outcome := rand.Intn(6) + 1",
|
||||||
|
];
|
||||||
|
|
||||||
|
// Run model
|
||||||
|
let output = sequence_classification_model.predict(input);
|
||||||
|
for label in output {
|
||||||
|
println!("{:?}", label);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Masked language model
|
||||||
|
let config = MaskedLanguageConfig::new(
|
||||||
|
ModelType::Roberta,
|
||||||
|
RemoteResource::from_pretrained(RobertaModelResources::CODEBERT_MLM),
|
||||||
|
RemoteResource::from_pretrained(RobertaConfigResources::CODEBERT_MLM),
|
||||||
|
RemoteResource::from_pretrained(RobertaVocabResources::CODEBERT_MLM),
|
||||||
|
Some(RemoteResource::from_pretrained(
|
||||||
|
RobertaMergesResources::CODEBERT_MLM,
|
||||||
|
)),
|
||||||
|
false,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(String::from("<mask>")),
|
||||||
|
);
|
||||||
|
|
||||||
|
let mask_language_model = MaskedLanguageModel::new(config)?;
|
||||||
|
// Define input
|
||||||
|
let input = [
|
||||||
|
"if (x is not None) <mask> (x>1)",
|
||||||
|
"<mask> x = if let <mask>(x_option) {}",
|
||||||
|
];
|
||||||
|
|
||||||
|
// Run model
|
||||||
|
let output = mask_language_model.predict(input)?;
|
||||||
|
for sentence_output in output {
|
||||||
|
println!("{:?}", sentence_output);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
@ -69,11 +69,21 @@ impl RobertaModelResources {
|
|||||||
"xlm-roberta-ner-es/model",
|
"xlm-roberta-ner-es/model",
|
||||||
"https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/rust_model.ot",
|
"https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/rust_model.ot",
|
||||||
);
|
);
|
||||||
/// Shared under Apache 2.0 licenseat <https://huggingface.co/sentence-transformers/all-distilroberta-v1>. Modified with conversion to C-array format.
|
/// Shared under Apache 2.0 license by the HuggingFace Inc. team at <https://huggingface.co/sentence-transformers/all-distilroberta-v1>. Modified with conversion to C-array format.
|
||||||
pub const ALL_DISTILROBERTA_V1: (&'static str, &'static str) = (
|
pub const ALL_DISTILROBERTA_V1: (&'static str, &'static str) = (
|
||||||
"all-distilroberta-v1/model",
|
"all-distilroberta-v1/model",
|
||||||
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/rust_model.ot",
|
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/rust_model.ot",
|
||||||
);
|
);
|
||||||
|
/// Shared under Apache 2.0 license by the HuggingFace Inc. team at <https://huggingface.co/huggingface/CodeBERTa-language-id>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERTA_LANGUAGE_ID: (&'static str, &'static str) = (
|
||||||
|
"codeberta-language-id/model",
|
||||||
|
"https://huggingface.co/huggingface/CodeBERTa-language-id/resolve/main/rust_model.ot",
|
||||||
|
);
|
||||||
|
/// Shared under MIT license by the Microsoft team at <https://github.com/microsoft/CodeBERT>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERT_MLM: (&'static str, &'static str) = (
|
||||||
|
"codebert-mlm/model",
|
||||||
|
"https://huggingface.co/microsoft/codebert-base-mlm/resolve/main/rust_model.ot",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RobertaConfigResources {
|
impl RobertaConfigResources {
|
||||||
@ -117,6 +127,16 @@ impl RobertaConfigResources {
|
|||||||
"all-distilroberta-v1/config",
|
"all-distilroberta-v1/config",
|
||||||
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/config.json",
|
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/config.json",
|
||||||
);
|
);
|
||||||
|
/// Shared under Apache 2.0 license by the HuggingFace Inc. team at <https://huggingface.co/huggingface/CodeBERTa-language-id>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERTA_LANGUAGE_ID: (&'static str, &'static str) = (
|
||||||
|
"codeberta-language-id/config",
|
||||||
|
"https://huggingface.co/huggingface/CodeBERTa-language-id/resolve/main/config.json",
|
||||||
|
);
|
||||||
|
/// Shared under MIT license by the Microsoft team at <https://github.com/microsoft/CodeBERT>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERT_MLM: (&'static str, &'static str) = (
|
||||||
|
"codebert-mlm/config",
|
||||||
|
"https://huggingface.co/microsoft/codebert-base-mlm/resolve/main/config.json",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RobertaVocabResources {
|
impl RobertaVocabResources {
|
||||||
@ -160,6 +180,16 @@ impl RobertaVocabResources {
|
|||||||
"all-distilroberta-v1/vocab",
|
"all-distilroberta-v1/vocab",
|
||||||
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/vocab.json",
|
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/vocab.json",
|
||||||
);
|
);
|
||||||
|
/// Shared under Apache 2.0 license by the HuggingFace Inc. team at <https://huggingface.co/huggingface/CodeBERTa-language-id>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERTA_LANGUAGE_ID: (&'static str, &'static str) = (
|
||||||
|
"codeberta-language-id/vocab",
|
||||||
|
"https://huggingface.co/huggingface/CodeBERTa-language-id/resolve/main/vocab.json",
|
||||||
|
);
|
||||||
|
/// Shared under MIT license by the Microsoft team at <https://github.com/microsoft/CodeBERT>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERT_MLM: (&'static str, &'static str) = (
|
||||||
|
"codebert-mlm/vocab",
|
||||||
|
"https://huggingface.co/microsoft/codebert-base-mlm/resolve/main/vocab.json",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RobertaMergesResources {
|
impl RobertaMergesResources {
|
||||||
@ -183,6 +213,16 @@ impl RobertaMergesResources {
|
|||||||
"all-distilroberta-v1/merges",
|
"all-distilroberta-v1/merges",
|
||||||
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/merges.txt",
|
"https://huggingface.co/sentence-transformers/all-distilroberta-v1/resolve/main/merges.txt",
|
||||||
);
|
);
|
||||||
|
/// Shared under Apache 2.0 license by the HuggingFace Inc. team at <https://huggingface.co/huggingface/CodeBERTa-language-id>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERTA_LANGUAGE_ID: (&'static str, &'static str) = (
|
||||||
|
"codeberta-language-id/merges",
|
||||||
|
"https://huggingface.co/huggingface/CodeBERTa-language-id/resolve/main/merges.txt",
|
||||||
|
);
|
||||||
|
/// Shared under MIT license by the Microsoft team at <https://github.com/microsoft/CodeBERT>. Modified with conversion to C-array format.
|
||||||
|
pub const CODEBERT_MLM: (&'static str, &'static str) = (
|
||||||
|
"codebert-mlm/merges",
|
||||||
|
"https://huggingface.co/microsoft/codebert-base-mlm/resolve/main/merges.txt",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct RobertaLMHead {
|
pub struct RobertaLMHead {
|
||||||
|
Loading…
Reference in New Issue
Block a user