[AutoNLP] Add english models for text classification (#4704)

* add english models

* add tests

Author: sijunhe

paddlenlp/experimental/autonlp/auto_trainer_base.py

@@ -67,9 +67,21 @@ def __init__(
         self.train_dataset = train_dataset
         self.eval_dataset = eval_dataset
         self.greater_is_better = greater_is_better
+        if language not in self.supported_languages:
+            raise ValueError(
+                f"'{language}' is not supported. Please choose among the following: {self.supported_languages}"
+            )
+
         self.language = language
         self.output_dir = output_dir
 
+    @property
+    @abstractmethod
+    def supported_languages(self) -> List[str]:
+        """
+        Override to store the supported languages for each auto trainer class
+        """
+
     @property
     @abstractmethod
     def _default_training_argument(self) -> TrainingArguments:

paddlenlp/experimental/autonlp/text_classification.py

@@ -89,6 +89,10 @@ def __init__(
                 f"'{problem_type}' is not a supported problem_type. Please select among ['multi_label', 'multi_class']"
             )
 
+    @property
+    def supported_languages(self) -> List[str]:
+        return ["Chinese", "English"]
+
     @property
     def _default_training_argument(self) -> TrainingArguments:
         return TrainingArguments(
@@ -129,19 +133,42 @@ def _model_candidates(self) -> List[Dict[str, Any]]:
                 "ernie-3.0-nano-zh",  # 4-layer, 312-hidden, 12-heads, 18M parameters.
             ],
         )
+        english_models = hp.choice(
+            "models",
+            [
+                # add deberta-v3 when we have it
+                "roberta-large",  # 24-layer, 1024-hidden, 16-heads, 334M parameters. Case-sensitive
+                "roberta-base",  # 12-layer, 768-hidden, 12-heads, 110M parameters. Case-sensitive
+                "distilroberta-base",  # 6-layer, 768-hidden, 12-heads, 66M parameters. Case-sensitive
+                "ernie-2.0-base-en",  # 12-layer, 768-hidden, 12-heads, 103M parameters. Trained on lower-cased English text.
+                "ernie-2.0-large-en",  # 24-layer, 1024-hidden, 16-heads, 336M parameters. Trained on lower-cased English text.
+                "distilbert-base-uncased",  # 6-layer, 768-hidden, 12-heads, 66M parameters
+            ],
+        )
         return [
             # fast learning: high LR, small early stop patience
             {
                 "preset": "finetune",
                 "language": "Chinese",
                 "trainer_type": "Trainer",
-                "EarlyStoppingCallback.early_stopping_patience": 2,
+                "EarlyStoppingCallback.early_stopping_patience": 5,
                 "TrainingArguments.per_device_train_batch_size": train_batch_size,
                 "TrainingArguments.per_device_eval_batch_size": train_batch_size * 2,
                 "TrainingArguments.num_train_epochs": 100,
                 "TrainingArguments.model_name_or_path": chinese_models,
                 "TrainingArguments.learning_rate": 3e-5,
             },
+            {
+                "preset": "finetune",
+                "language": "English",
+                "trainer_type": "Trainer",
+                "EarlyStoppingCallback.early_stopping_patience": 5,
+                "TrainingArguments.per_device_train_batch_size": train_batch_size,
+                "TrainingArguments.per_device_eval_batch_size": train_batch_size * 2,
+                "TrainingArguments.num_train_epochs": 100,
+                "TrainingArguments.model_name_or_path": english_models,
+                "TrainingArguments.learning_rate": 3e-5,
+            },
             # slow learning: small LR, large early stop patience
             {
                 "preset": "finetune",
@@ -154,6 +181,17 @@ def _model_candidates(self) -> List[Dict[str, Any]]:
                 "TrainingArguments.model_name_or_path": chinese_models,
                 "TrainingArguments.learning_rate": 5e-6,
             },
+            {
+                "preset": "finetune",
+                "language": "English",
+                "trainer_type": "Trainer",
+                "EarlyStoppingCallback.early_stopping_patience": 5,
+                "TrainingArguments.per_device_train_batch_size": train_batch_size,
+                "TrainingArguments.per_device_eval_batch_size": train_batch_size * 2,
+                "TrainingArguments.num_train_epochs": 100,
+                "TrainingArguments.model_name_or_path": english_models,
+                "TrainingArguments.learning_rate": 5e-6,
+            },
             # Note: prompt tuning candidates not included for now due to lack of inference capability
         ]
 

tests/experimental/autonlp/test_text_classification.py

@@ -268,6 +268,37 @@ def test_untrained_auto_trainer(self):
                 # test export
                 auto_trainer.export(temp_dir)
 
+    def test_unsupported_languages(self):
+        with TemporaryDirectory() as temp_dir:
+            train_ds = copy.deepcopy(self.multi_class_train_ds)
+            dev_ds = copy.deepcopy(self.multi_class_dev_ds)
+            with self.assertRaises(ValueError):
+                AutoTrainerForTextClassification(
+                    train_dataset=train_ds,
+                    eval_dataset=dev_ds,
+                    label_column="label_desc",
+                    text_column="sentence",
+                    language="Spanish",  # spanish is unsupported for now
+                    output_dir=temp_dir,
+                )
+
+    def test_model_language_filter(self):
+        with TemporaryDirectory() as temp_dir:
+            train_ds = copy.deepcopy(self.multi_class_train_ds)
+            dev_ds = copy.deepcopy(self.multi_class_dev_ds)
+            auto_trainer = AutoTrainerForTextClassification(
+                train_dataset=train_ds,
+                eval_dataset=dev_ds,
+                label_column="label_desc",
+                text_column="sentence",
+                language="Chinese",
+                output_dir=temp_dir,
+            )
+            for language in auto_trainer.supported_languages:
+                model_candidates = auto_trainer._filter_model_candidates(language=language)
+                for candidate in model_candidates:
+                    self.assertEqual(candidate["language"], language)
+
 
 if __name__ == "__main__":
     unittest.main()