make ZeroShotClassification support the roberta-large-mnli

I don't know if it's correct, so a review would be appreciated
guillaume-be · Sep 21, 2023 · 49b7a5d · 49b7a5d
1 parent 0a1c3df
commit 49b7a5d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/pipelines/zero_shot_classification.rs b/src/pipelines/zero_shot_classification.rs
@@ -322,7 +322,7 @@ 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)?,
                     ))
@@ -491,7 +491,7 @@ impl ZeroShotClassificationOption {
                     .forward_t(
                         input_ids,
                         mask,
-                        token_type_ids,
+                        None,
                         position_ids,
                         input_embeds,
                         train,