Fix for KeytoValue transformer (#4866)

* fix for keytovalue

* fix for keytovalue

* resolving comments

Author: Lynx1820

src/Microsoft.ML.Data/Transforms/KeyToValue.cs

@@ -568,9 +568,9 @@ public void SaveAsOnnx(OnnxContext ctx)
 
                     if (!ctx.ContainsColumn(inputColumnName))
                         continue;
-
+                    string srcVariableName = ctx.GetVariableName(inputColumnName);
                     var dstVariableName = ctx.AddIntermediateVariable(_types[iinfo], info.outputColumnName, true);
-                    if (!_kvMaps[iinfo].SaveOnnx(ctx, inputColumnName, dstVariableName))
+                    if (!_kvMaps[iinfo].SaveOnnx(ctx, srcVariableName, dstVariableName))
                     {
                         ctx.RemoveColumn(inputColumnName, true);
                     }

test/Microsoft.ML.Tests/OnnxConversionTest.cs

@@ -9,6 +9,7 @@
 using System.Runtime.InteropServices;
 using System.Text.RegularExpressions;
 using Google.Protobuf;
+using Google.Protobuf.WellKnownTypes;
 using Microsoft.ML.Data;
 using Microsoft.ML.EntryPoints;
 using Microsoft.ML.Model.OnnxConverter;
@@ -1150,7 +1151,7 @@ public void IndicateMissingValuesOnnxConversionTest()
         [InlineData(DataKind.Int64)]
         [InlineData(DataKind.Double)]
         [InlineData(DataKind.String)]
-        public void ValueToKeyandKeyToValueMappingOnnxConversionTest(DataKind valueType)
+        public void ValueToKeyMappingOnnxConversionTest(DataKind valueType)
         {
             var mlContext = new MLContext(seed: 1);
             string filePath = GetDataPath("type-conversion.txt");
@@ -1160,9 +1161,8 @@ public void ValueToKeyandKeyToValueMappingOnnxConversionTest(DataKind valueType)
                 new TextLoader.Column("Value", valueType, 0, 0)
             };
             var dataView = mlContext.Data.LoadFromTextFile(filePath, columns);
+            var pipeline = mlContext.Transforms.Conversion.MapValueToKey("Key", "Value");
 
-            var pipeline = mlContext.Transforms.Conversion.MapValueToKey("Key", "Value").
-                Append(mlContext.Transforms.Conversion.MapKeyToValue("ValueOutput", "Key"));
             var model = pipeline.Fit(dataView);
             var mlnetResult = model.Transform(dataView);
 
@@ -1176,13 +1176,55 @@ public void ValueToKeyandKeyToValueMappingOnnxConversionTest(DataKind valueType)
                 var onnxEstimator = mlContext.Transforms.ApplyOnnxModel(onnxModelPath);
                 var onnxTransformer = onnxEstimator.Fit(dataView);
                 var onnxResult = onnxTransformer.Transform(dataView);
-
-                CompareResults("ValueOutput", "ValueOutput", mlnetResult, onnxResult);
                 CompareSelectedColumns<uint>("Key", "Key", mlnetResult, onnxResult);
             }
             Done();
         }
 
+        [Theory]
+        [InlineData(DataKind.Single)]
+        [InlineData(DataKind.Int64)]
+        [InlineData(DataKind.Double)]
+        [InlineData(DataKind.String)]
+        public void KeyToValueMappingOnnxConversionTest(DataKind valueType)
+        {
+            var mlContext = new MLContext(seed: 1);
+            string filePath = GetDataPath("type-conversion.txt");
+
+            TextLoader.Column[] columns = new[]
+            {
+                new TextLoader.Column("Value", valueType, 0, 0)
+            };
+            var dataView = mlContext.Data.LoadFromTextFile(filePath, columns);
+            IEstimator<ITransformer>[] pipelines =
+            {
+                mlContext.Transforms.Conversion.MapValueToKey("Key", "Value").
+                Append(mlContext.Transforms.Conversion.MapKeyToValue("Value", "Key")),
+
+                mlContext.Transforms.Conversion.MapValueToKey("Value").
+                Append(mlContext.Transforms.Conversion.MapKeyToValue("Value"))
+            };
+            for (int i = 0; i < pipelines.Length; i++)
+            {
+                var model = pipelines[i].Fit(dataView);
+                var mlnetResult = model.Transform(dataView);
+
+                var onnxModel = mlContext.Model.ConvertToOnnxProtobuf(model, dataView);
+                var onnxFileName = "KeyToValue.onnx";
+                var onnxModelPath = GetOutputPath(onnxFileName);
+                SaveOnnxModel(onnxModel, onnxModelPath, null);
+
+                if (IsOnnxRuntimeSupported())
+                {
+                    var onnxEstimator = mlContext.Transforms.ApplyOnnxModel(onnxModelPath);
+                    var onnxTransformer = onnxEstimator.Fit(dataView);
+                    var onnxResult = onnxTransformer.Transform(dataView);
+                    CompareResults("Value", "Value", mlnetResult, onnxResult);
+                }
+            }
+            Done();
+        }
+
         private class TextData
         {
             public string Text { get; set; }
@@ -1355,40 +1397,6 @@ public void OptionalColumnOnnxTest(DataKind dataKind)
             Done();
         }
 
-        [Fact]
-        public void KeyToValueOnnxConversionTest()
-        {
-            var mlContext = new MLContext(seed: 1);
-
-            string dataPath = GetDataPath("breast-cancer.txt");
-            var dataView = mlContext.Data.LoadFromTextFile<BreastCancerMulticlassExample>(dataPath,
-                separatorChar: '\t',
-                hasHeader: true);
-
-            var pipeline = mlContext.Transforms.Conversion.MapValueToKey("LabelKey", "Label").
-                Append(mlContext.Transforms.Conversion.MapKeyToValue("LabelValue", "LabelKey"));
-
-            var model = pipeline.Fit(dataView);
-            var transformedData = model.Transform(dataView);
-            var onnxModel = mlContext.Model.ConvertToOnnxProtobuf(model, dataView);
-
-            var onnxFileName = "KeyToValue.onnx";
-            var onnxModelPath = GetOutputPath(onnxFileName);
-
-            SaveOnnxModel(onnxModel, onnxModelPath, null);
-
-            if (IsOnnxRuntimeSupported())
-            {
-                // Evaluate the saved ONNX model using the data used to train the ML.NET pipeline.
-                var onnxEstimator = mlContext.Transforms.ApplyOnnxModel(onnxModelPath);
-                var onnxTransformer = onnxEstimator.Fit(dataView);
-                var onnxResult = onnxTransformer.Transform(dataView);
-                CompareSelectedColumns<ReadOnlyMemory<Char>>("LabelValue", "LabelValue", transformedData, onnxResult);
-            }
-
-            Done();
-        }
-
         [Fact]
         public void MulticlassTrainersOnnxConversionTest()
         {