Add sweepable estimator to NER (#6965)

dotnet · Jan 19, 2024 · 125b6d5 · 125b6d5
1 parent 48b6fbe
commit 125b6d5
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 3 deletions.
diff --git a/src/Microsoft.ML.AutoML/CodeGen/estimator-schema.json b/src/Microsoft.ML.AutoML/CodeGen/estimator-schema.json
@@ -74,7 +74,8 @@
               "TextClassifcation",
               "SentenceSimilarity",
               "ObjectDetection",
-              "QuestionAnswering"
+              "QuestionAnswering",
+              "NamedEntityRecognition"
             ]
           },
           "nugetDependencies": {

diff --git a/src/Microsoft.ML.AutoML/CodeGen/named_entity_recognition_search_space.json b/src/Microsoft.ML.AutoML/CodeGen/named_entity_recognition_search_space.json
@@ -0,0 +1,36 @@
+{
+  "$schema": "./search-space-schema.json#",
+  "name": "named_entity_recognition_option",
+  "search_space": [
+    {
+      "name": "PredictionColumnName",
+      "type": "string",
+      "default": "predictedLabel"
+    },
+    {
+      "name": "LabelColumnName",
+      "type": "string",
+      "default": "Label"
+    },
+    {
+      "name": "Sentence1ColumnName",
+      "type": "string",
+      "default": "Sentence"
+    },
+    {
+      "name": "BatchSize",
+      "type": "integer",
+      "default": 32
+    },
+    {
+      "name": "MaxEpochs",
+      "type": "integer",
+      "default": 10
+    },
+    {
+      "name": "Architecture",
+      "type": "bertArchitecture",
+      "default": "BertArchitecture.Roberta"
+    }
+  ]
+}
diff --git a/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json b/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json
@@ -167,7 +167,8 @@
         "text_classification_option",
         "sentence_similarity_option",
         "object_detection_option",
-        "question_answering_option"
+        "question_answering_option",
+        "named_entity_recognition_option"
       ]
     },
     "option_name": {
@@ -238,7 +239,8 @@
         "AnswerIndexStartColumnName",
         "predictedAnswerColumnName",
         "TopKAnswers",
-        "TargetType"
+        "TargetType",
+        "PredictionColumnName"
       ]
     },
     "option_type": {

diff --git a/src/Microsoft.ML.AutoML/CodeGen/trainer-estimators.json b/src/Microsoft.ML.AutoML/CodeGen/trainer-estimators.json
@@ -539,6 +539,13 @@
       "usingStatements": [ "Microsoft.ML", "Microsoft.ML.Trainers", "Microsoft.ML.TorchSharp" ],
       "searchOption": "question_answering_option"
     },
+    {
+      "functionName": "NamedEntityRecognition",
+      "estimatorTypes": [ "MultiClassification" ],
+      "nugetDependencies": [ "Microsoft.ML", "Microsoft.ML.TorchSharp" ],
+      "usingStatements": [ "Microsoft.ML", "Microsoft.ML.Trainers", "Microsoft.ML.TorchSharp" ],
+      "searchOption": "named_entity_recognition_option"
+    },
     {
       "functionName": "ForecastBySsa",
       "estimatorTypes": [ "Forecasting" ],

diff --git a/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/NamedEntityRecognitionMulti.cs b/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/NamedEntityRecognitionMulti.cs
@@ -0,0 +1,26 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Microsoft.ML.TorchSharp;
+using Microsoft.ML.TorchSharp.NasBert;
+
+namespace Microsoft.ML.AutoML.CodeGen
+{
+    internal partial class NamedEntityRecognitionMulti
+    {
+        public override IEstimator<ITransformer> BuildFromOption(MLContext context, NamedEntityRecognitionOption param)
+        {
+            return context.MulticlassClassification.Trainers.NamedEntityRecognition(
+                   labelColumnName: param.LabelColumnName,
+                   outputColumnName: param.PredictionColumnName,
+                   sentence1ColumnName: param.Sentence1ColumnName,
+                   batchSize: param.BatchSize,
+                   maxEpochs: param.MaxEpochs,
+                   architecture: BertArchitecture.Roberta);
+        }
+    }
+}