Revert "Add type resolving for ONNX models"

2 files changed, 6 insertions, 6 deletions

diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RankProfileList.java b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RankProfileList.java
index 22a32c8fd65..84442fedc48 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RankProfileList.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RankProfileList.java
@@ -126,8 +126,8 @@ public class RankProfileList extends Derived implements RankProfilesConfig.Produ
                 OnnxModelsConfig.Model.Builder modelBuilder = new OnnxModelsConfig.Model.Builder();
                 modelBuilder.name(model.getName());
                 modelBuilder.fileref(model.getFileReference());
-                model.getInputMap().forEach((name, source) -> modelBuilder.input(new OnnxModelsConfig.Model.Input.Builder().name(name).source(source)));
-                model.getOutputMap().forEach((name, as) -> modelBuilder.output(new OnnxModelsConfig.Model.Output.Builder().name(name).as(as)));
+                model.getInputMap().forEach(mapper -> modelBuilder.input(new OnnxModelsConfig.Model.Input.Builder().name(mapper.getOnnxName()).source(mapper.getVespaName())));
+                model.getOutputMap().forEach(mapper -> modelBuilder.output(new OnnxModelsConfig.Model.Output.Builder().name(mapper.getOnnxName()).as(mapper.getVespaName())));
                 builder.model(modelBuilder);
             }
         }
diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
index 56a5d539906..87eaaf0387a 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
@@ -448,10 +448,10 @@ public class RawRankProfile implements RankProfilesConfig.Producer {
             Set<String> functionNames = rankProfile.getFunctions().keySet();
             if (functionNames.isEmpty()) return;
             for (OnnxModel onnxModel: rankProfile.getSearch().onnxModels().asMap().values()) {
-                for (Map.Entry<String, String> mapping : onnxModel.getInputMap().entrySet()) {
-                    String source = mapping.getValue();
+                for (OnnxModel.OnnxNameMapping mapping : onnxModel.getInputMap()) {
+                    String source = mapping.getVespaName();
                     if (functionNames.contains(source)) {
-                        onnxModel.addInputNameMapping(mapping.getKey(), "rankingExpression(" + source + ")");
+                        mapping.setVespaName("rankingExpression(" + source + ")");
                     }
                 }
             }
@@ -462,7 +462,7 @@ public class RawRankProfile implements RankProfilesConfig.Producer {
             Set<ReferenceNode> replacedSummaryFeatures = new HashSet<>();
             for (Iterator<ReferenceNode> i = summaryFeatures.iterator(); i.hasNext(); ) {
                 ReferenceNode referenceNode = i.next();
-                ReferenceNode replacedNode = (ReferenceNode) OnnxModelTransformer.transformFeature(referenceNode, rankProfile);
+                ReferenceNode replacedNode = OnnxModelTransformer.transformFeature(referenceNode, rankProfile.getSearch());
                 if (referenceNode != replacedNode) {
                     replacedSummaryFeatures.add(replacedNode);
                     i.remove();