make ZeroShotClassification support the roberta-large-mnli (#422)

* make ZeroShotClassification support the roberta-large-mnli I don't know if it's correct, so a review would be appreciated * change error message for ZeroShotClassification and Roberta models * Revert token type id update --------- Co-authored-by: Charles Samuels <[email protected]> Co-authored-by: guillaume-be <[email protected]>
guillaume-be · Oct 21, 2023 · ce48aa0 · ce48aa0
1 parent 7a5c42d
commit ce48aa0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -75,7 +75,7 @@ hf-tokenizers = ["tokenizers"]
 features = ["doc-only"]
 
 [dependencies]
-rust_tokenizers = "8.1"
+rust_tokenizers = "8.1.1"
 tch = "0.13.0"
 serde_json = "1"
 serde = { version = "1", features = ["derive"] }

diff --git a/src/pipelines/zero_shot_classification.rs b/src/pipelines/zero_shot_classification.rs
@@ -322,13 +322,13 @@ impl ZeroShotClassificationOption {
                 }
             }
             ModelType::Roberta => {
-                if let ConfigOption::Bert(config) = model_config {
+                if let ConfigOption::Roberta(config) = model_config {
                     Ok(Self::Roberta(
                         RobertaForSequenceClassification::new(var_store.root(), config)?,
                     ))
                 } else {
                     Err(RustBertError::InvalidConfigurationError(
-                        "You can only supply a BertConfig for Roberta!".to_string(),
+                        "You can only supply a RobertaConfig for Roberta!".to_string(),
                     ))
                 }
             }