diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_job.py
index 1f6a31144c29..63e468c4f2e5 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_job.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_job.py
@@ -8,10 +8,10 @@
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import AutoMLJob as RestAutoMLJob
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase, TaskType
-from azure.ai.ml._restclient.v2023_04_01_preview.models import TextClassification as RestTextClassification
 from azure.ai.ml._restclient.v2023_04_01_preview.models._azure_machine_learning_workspaces_enums import (
     ClassificationPrimaryMetrics,
 )
+from azure.ai.ml._restclient.v2024_01_01_preview.models import TextClassification as RestTextClassification
 from azure.ai.ml._utils.utils import camel_to_snake, is_data_binding_expression
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml.constants._job.automl import AutoMLConstants
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_multilabel_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_multilabel_job.py
index 2f9dfedfff1c..d277a1657e05 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_multilabel_job.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_multilabel_job.py
@@ -8,7 +8,7 @@
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import AutoMLJob as RestAutoMLJob
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ClassificationMultilabelPrimaryMetrics, JobBase, TaskType
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.v2024_01_01_preview.models import (
     TextClassificationMultilabel as RestTextClassificationMultilabel,
 )
 from azure.ai.ml._utils.utils import camel_to_snake, is_data_binding_expression
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_ner_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_ner_job.py
index 2dc2f7100b3a..38cd9b3d6416 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_ner_job.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_ner_job.py
@@ -8,10 +8,10 @@
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import AutoMLJob as RestAutoMLJob
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase, TaskType
-from azure.ai.ml._restclient.v2023_04_01_preview.models import TextNer as RestTextNER
 from azure.ai.ml._restclient.v2023_04_01_preview.models._azure_machine_learning_workspaces_enums import (
     ClassificationPrimaryMetrics,
 )
+from azure.ai.ml._restclient.v2024_01_01_preview.models import TextNer as RestTextNER
 from azure.ai.ml._utils.utils import camel_to_snake, is_data_binding_expression
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml.constants._job.automl import AutoMLConstants
diff --git a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_automl_nlp_schema.py b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_automl_nlp_schema.py
index 90af08f318ca..88b867317276 100644
--- a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_automl_nlp_schema.py
+++ b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_automl_nlp_schema.py
@@ -19,7 +19,6 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models._models_py3 import (
     NlpVerticalLimitSettings as RestNlpVerticalLimitSettings,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models._models_py3 import TextClassification as RestTextClassification
 from azure.ai.ml._restclient.v2023_04_01_preview.models._models_py3 import (
     TextClassificationMultilabel as RestTextClassificationMultilabel,
 )
@@ -30,6 +29,7 @@
     LogVerbosity,
     MLTableJobInput,
 )
+from azure.ai.ml._restclient.v2024_01_01_preview.models._models_py3 import TextClassification as RestTextClassification
 from azure.ai.ml._scope_dependent_operations import OperationScope
 from azure.ai.ml._utils.utils import dump_yaml_to_file, load_yaml, to_iso_duration_format_mins
 from azure.ai.ml.automl import (
diff --git a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_job.py b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_job.py
index 57cf6badb6f6..395331cdb3e0 100644
--- a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_job.py
+++ b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_job.py
@@ -12,12 +12,11 @@
     NlpVerticalFeaturizationSettings,
     NlpVerticalLimitSettings,
     SamplingAlgorithmType,
-    TextClassification,
 )
 from azure.ai.ml._restclient.v2023_04_01_preview.models._azure_machine_learning_workspaces_enums import (
     ClassificationPrimaryMetrics,
 )
-from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput
+from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput, TextClassification
 from azure.ai.ml._utils.utils import to_iso_duration_format_mins
 from azure.ai.ml.automl import text_classification
 from azure.ai.ml.constants._common import AssetTypes
diff --git a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_multilabel_job.py b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_multilabel_job.py
index f0f3e1bac816..7abe4a8183f5 100644
--- a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_multilabel_job.py
+++ b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_classification_multilabel_job.py
@@ -14,13 +14,12 @@
     NlpVerticalFeaturizationSettings,
     NlpVerticalLimitSettings,
     SamplingAlgorithmType,
-    TextClassificationMultilabel,
 )
 from azure.ai.ml._restclient.v2023_04_01_preview.models import UserIdentity as RestUserIdentity
 from azure.ai.ml._restclient.v2023_04_01_preview.models._azure_machine_learning_workspaces_enums import (
     ClassificationPrimaryMetrics,
 )
-from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput
+from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput, TextClassificationMultilabel
 from azure.ai.ml._utils.utils import to_iso_duration_format_mins
 from azure.ai.ml.automl import text_classification_multilabel
 from azure.ai.ml.constants._common import AssetTypes
diff --git a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_ner_job.py b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_ner_job.py
index 2a9a608f7bee..ae626a93fdd6 100644
--- a/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_ner_job.py
+++ b/sdk/ml/azure-ai-ml/tests/automl_job/unittests/test_text_ner_job.py
@@ -13,13 +13,12 @@
     NlpVerticalLimitSettings,
     SamplingAlgorithmType,
     TaskType,
-    TextNer,
 )
 from azure.ai.ml._restclient.v2023_04_01_preview.models import UserIdentity as RestUserIdentity
 from azure.ai.ml._restclient.v2023_04_01_preview.models._azure_machine_learning_workspaces_enums import (
     ClassificationPrimaryMetrics,
 )
-from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput
+from azure.ai.ml._restclient.v2024_01_01_preview.models import MLTableJobInput, TextNer
 from azure.ai.ml._utils.utils import to_iso_duration_format_mins
 from azure.ai.ml.automl import text_ner
 from azure.ai.ml.constants._common import AssetTypes