From 3eeaf352d1baf50ae40dd85327472f07ea58f32b Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 24 Sep 2021 17:34:21 -0400
Subject: [PATCH 1/8] Stubbing out onnx provenance serialization.

---
 Core/pom.xml                                  |   4 +
 .../java/org/tribuo/onnx/ONNXExportable.java  |  23 +++
 Interop/ONNX/pom.xml                          |   5 +
 .../interop/onnx/ONNXExternalModel.java       | 162 ++++++++++++------
 pom.xml                                       |  11 +-
 5 files changed, 147 insertions(+), 58 deletions(-)

diff --git a/Core/pom.xml b/Core/pom.xml
index 4d03ecffd..c047464d1 100644
--- a/Core/pom.xml
+++ b/Core/pom.xml
@@ -37,6 +37,10 @@
             <groupId>com.oracle.labs.olcut</groupId>
             <artifactId>olcut-core</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.oracle.labs.olcut</groupId>
+            <artifactId>olcut-config-protobuf</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.google.protobuf</groupId>
             <artifactId>protobuf-java</artifactId>
diff --git a/Core/src/main/java/org/tribuo/onnx/ONNXExportable.java b/Core/src/main/java/org/tribuo/onnx/ONNXExportable.java
index 0c1c54ebe..b10978426 100644
--- a/Core/src/main/java/org/tribuo/onnx/ONNXExportable.java
+++ b/Core/src/main/java/org/tribuo/onnx/ONNXExportable.java
@@ -17,6 +17,9 @@
 package org.tribuo.onnx;
 
 import ai.onnx.proto.OnnxMl;
+import com.oracle.labs.mlrg.olcut.config.protobuf.ProtoProvenanceSerialization;
+import com.oracle.labs.mlrg.olcut.provenance.io.ProvenanceSerialization;
+import org.tribuo.provenance.ModelProvenance;
 
 import java.io.BufferedOutputStream;
 import java.io.FileOutputStream;
@@ -29,6 +32,17 @@
  */
 public interface ONNXExportable {
 
+    /**
+     * The provenance serializer.
+     */
+    public static final ProvenanceSerialization SERIALIZER = new ProtoProvenanceSerialization(true);
+
+    /**
+     * The name of the ONNX metadata field where the provenance information is stored
+     * in exported models.
+     */
+    public static final String PROVENANCE_METADATA_FIELD = "TRIBUO_PROVENANCE";
+
     /**
      * Exports this {@link org.tribuo.Model} as an ONNX protobuf.
      * @param domain A reverse-DNS name to namespace the model (e.g., org.tribuo.classification.sgd.linear).
@@ -74,4 +88,13 @@ default public void saveONNXModel(String domain, long modelVersion, Path outputP
         }
     }
 
+    /**
+     * Serializes the model provenance to a String.
+     * @param provenance The provenance to serialize.
+     * @return The serialized form of the ModelProvenance.
+     */
+    default public String serializeProvenance(ModelProvenance provenance) {
+        return SERIALIZER.marshalAndSerialize(provenance);
+    }
+
 }
diff --git a/Interop/ONNX/pom.xml b/Interop/ONNX/pom.xml
index f80b135e0..f69d4d472 100644
--- a/Interop/ONNX/pom.xml
+++ b/Interop/ONNX/pom.xml
@@ -67,6 +67,11 @@
             <artifactId>tribuo-util-tokenization</artifactId>
             <version>${project.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.oracle.labs.olcut</groupId>
+            <artifactId>olcut-config-protobuf</artifactId>
+            <version>${olcut.version}</version>
+        </dependency>
         <dependency>
             <groupId>com.microsoft.onnxruntime</groupId>
             <artifactId>onnxruntime</artifactId>
diff --git a/Interop/ONNX/src/main/java/org/tribuo/interop/onnx/ONNXExternalModel.java b/Interop/ONNX/src/main/java/org/tribuo/interop/onnx/ONNXExternalModel.java
index a5dd3ed03..d5a77cd9c 100644
--- a/Interop/ONNX/src/main/java/org/tribuo/interop/onnx/ONNXExternalModel.java
+++ b/Interop/ONNX/src/main/java/org/tribuo/interop/onnx/ONNXExternalModel.java
@@ -22,7 +22,10 @@
 import ai.onnxruntime.OrtEnvironment;
 import ai.onnxruntime.OrtException;
 import ai.onnxruntime.OrtSession;
+import com.oracle.labs.mlrg.olcut.config.protobuf.ProtoProvenanceSerialization;
 import com.oracle.labs.mlrg.olcut.provenance.Provenance;
+import com.oracle.labs.mlrg.olcut.provenance.io.ProvenanceSerialization;
+import com.oracle.labs.mlrg.olcut.provenance.io.ProvenanceSerializationException;
 import com.oracle.labs.mlrg.olcut.provenance.primitives.LongProvenance;
 import com.oracle.labs.mlrg.olcut.provenance.primitives.StringProvenance;
 import com.oracle.labs.mlrg.olcut.util.Pair;
@@ -37,6 +40,7 @@
 import org.tribuo.interop.ExternalModel;
 import org.tribuo.interop.ExternalTrainerProvenance;
 import org.tribuo.math.la.SparseVector;
+import org.tribuo.onnx.ONNXExportable;
 import org.tribuo.provenance.DatasetProvenance;
 import org.tribuo.provenance.ModelProvenance;
 
@@ -52,6 +56,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -90,8 +95,8 @@ private ONNXExternalModel(String name, ModelProvenance provenance,
         this.inputName = inputName;
         this.featureTransformer = featureTransformer;
         this.outputTransformer = outputTransformer;
-        this.env = OrtEnvironment.getEnvironment("tribuo-"+name);
-        this.session = env.createSession(modelArray,options);
+        this.env = OrtEnvironment.getEnvironment("tribuo-" + name);
+        this.session = env.createSession(modelArray, options);
     }
 
     private ONNXExternalModel(String name, ModelProvenance provenance,
@@ -99,21 +104,22 @@ private ONNXExternalModel(String name, ModelProvenance provenance,
                               int[] featureForwardMapping, int[] featureBackwardMapping,
                               byte[] modelArray, OrtSession.SessionOptions options, String inputName,
                               ExampleTransformer featureTransformer, OutputTransformer<T> outputTransformer) throws OrtException {
-        super(name,provenance,featureIDMap,outputIDInfo,featureForwardMapping,featureBackwardMapping,
+        super(name, provenance, featureIDMap, outputIDInfo, featureForwardMapping, featureBackwardMapping,
                 outputTransformer.generatesProbabilities());
         this.modelArray = modelArray;
         this.options = options;
         this.inputName = inputName;
         this.featureTransformer = featureTransformer;
         this.outputTransformer = outputTransformer;
-        this.env = OrtEnvironment.getEnvironment("tribuo-"+name);
-        this.session = env.createSession(modelArray,options);
+        this.env = OrtEnvironment.getEnvironment("tribuo-" + name);
+        this.session = env.createSession(modelArray, options);
     }
 
     /**
      * Closes the session and rebuilds it using the supplied options.
      * <p>
      * Used to select a different backend, or change the number of inference threads etc.
+     *
      * @param newOptions The new session options.
      * @throws OrtException If the model failed to rebuild the session with the supplied options.
      */
@@ -123,7 +129,7 @@ public synchronized void rebuild(OrtSession.SessionOptions newOptions) throws Or
             options.close();
         }
         options = newOptions;
-        env.createSession(modelArray,newOptions);
+        env.createSession(modelArray, newOptions);
     }
 
     @Override
@@ -131,7 +137,7 @@ protected OnnxTensor convertFeatures(SparseVector input) {
         try {
             return featureTransformer.transform(env, input);
         } catch (OrtException e) {
-            throw new IllegalStateException("Failed to construct input OnnxTensor",e);
+            throw new IllegalStateException("Failed to construct input OnnxTensor", e);
         }
     }
 
@@ -140,7 +146,7 @@ protected OnnxTensor convertFeaturesList(List<SparseVector> input) {
         try {
             return featureTransformer.transform(env, input);
         } catch (OrtException e) {
-            throw new IllegalStateException("Failed to construct input OnnxTensor",e);
+            throw new IllegalStateException("Failed to construct input OnnxTensor", e);
         }
     }
 
@@ -148,6 +154,7 @@ protected OnnxTensor convertFeaturesList(List<SparseVector> input) {
      * Runs the session to make a prediction.
      * <p>
      * Closes the input tensor after the prediction has been made.
+     *
      * @param input The input in the external model's format.
      * @return A tensor representing the output.
      */
@@ -155,29 +162,30 @@ protected OnnxTensor convertFeaturesList(List<SparseVector> input) {
     protected List<OnnxValue> externalPrediction(OnnxTensor input) {
         try {
             // Note the output of the session is closed by the conversion methods, and should not be closed by the result object.
-            OrtSession.Result output = session.run(Collections.singletonMap(inputName,input));
+            OrtSession.Result output = session.run(Collections.singletonMap(inputName, input));
             input.close();
             ArrayList<OnnxValue> outputs = new ArrayList<>();
-            for (Map.Entry<String,OnnxValue> v : output) {
+            for (Map.Entry<String, OnnxValue> v : output) {
                 outputs.add(v.getValue());
             }
             return outputs;
         } catch (OrtException e) {
-            throw new IllegalStateException("Failed to execute ONNX model",e);
+            throw new IllegalStateException("Failed to execute ONNX model", e);
         }
     }
 
     /**
      * Converts a tensor into a prediction.
      * Closes the output tensor after it's been converted.
-     * @param output The output of the external model.
+     *
+     * @param output           The output of the external model.
      * @param numValidFeatures The number of valid features in the input.
-     * @param example The input example, used to construct the Prediction.
+     * @param example          The input example, used to construct the Prediction.
      * @return A {@link Prediction} representing this tensor output.
      */
     @Override
     protected Prediction<T> convertOutput(List<OnnxValue> output, int numValidFeatures, Example<T> example) {
-        Prediction<T> pred = outputTransformer.transformToPrediction(output,outputIDInfo,numValidFeatures,example);
+        Prediction<T> pred = outputTransformer.transformToPrediction(output, outputIDInfo, numValidFeatures, example);
         OnnxValue.close(output);
         return pred;
     }
@@ -185,14 +193,15 @@ protected Prediction<T> convertOutput(List<OnnxValue> output, int numValidFeatur
     /**
      * Converts a tensor into a prediction.
      * Closes the output tensor after it's been converted.
-     * @param output The output of the external model.
+     *
+     * @param output           The output of the external model.
      * @param numValidFeatures An array with the number of valid features in each example.
-     * @param examples The input examples, used to construct the Predictions.
+     * @param examples         The input examples, used to construct the Predictions.
      * @return A list of {@link Prediction} representing this tensor output.
      */
     @Override
     protected List<Prediction<T>> convertOutput(List<OnnxValue> output, int[] numValidFeatures, List<Example<T>> examples) {
-        List<Prediction<T>> predictions = outputTransformer.transformToBatchPrediction(output,outputIDInfo,numValidFeatures,examples);
+        List<Prediction<T>> predictions = outputTransformer.transformToBatchPrediction(output, outputIDInfo, numValidFeatures, examples);
         OnnxValue.close(output);
         return predictions;
     }
@@ -204,13 +213,13 @@ public Map<String, List<Pair<String, Double>>> getTopFeatures(int n) {
 
     @Override
     protected synchronized Model<T> copy(String newName, ModelProvenance newProvenance) {
-        byte[] newModelArray = Arrays.copyOf(modelArray,modelArray.length);
+        byte[] newModelArray = Arrays.copyOf(modelArray, modelArray.length);
         try {
             return new ONNXExternalModel<>(newName, newProvenance, featureIDMap, outputIDInfo,
                     featureForwardMapping, featureBackwardMapping,
                     newModelArray, options, inputName, featureTransformer, outputTransformer);
         } catch (OrtException e) {
-            throw new IllegalStateException("Failed to copy ONNX model",e);
+            throw new IllegalStateException("Failed to copy ONNX model", e);
         }
     }
 
@@ -220,7 +229,7 @@ public void close() {
             try {
                 session.close();
             } catch (OrtException e) {
-                logger.log(Level.SEVERE,"Exception thrown when closing session",e);
+                logger.log(Level.SEVERE, "Exception thrown when closing session", e);
             }
         }
         if (options != null) {
@@ -230,54 +239,97 @@ public void close() {
             try {
                 env.close();
             } catch (OrtException e) {
-                logger.log(Level.SEVERE,"Exception thrown when closing environment",e);
+                logger.log(Level.SEVERE, "Exception thrown when closing environment", e);
             }
         }
     }
 
+    /**
+     * Returns the model provenance from the ONNX model if that
+     * model was trained in Tribuo.
+     * <p>
+     * Tribuo's ONNX export functionality stores the model provenance inside the
+     * ONNX file in the metadata field {@link ONNXExportable#PROVENANCE_METADATA_FIELD},
+     * and this method provides the access point for it.
+     * <p>
+     * Note it is different from the {@link Model#getProvenance()} call which
+     * returns information about the ONNX file itself, and when the {@code ONNXExternalModel}
+     * was created. It does not replace that provenance because instantiating this provenance
+     * may require classes which are not present on the classpath at deployment time.
+     *
+     * @return The model provenance from the original Tribuo training run, if it exists, and
+     * returns {@link Optional#empty()} otherwise.
+     */
+    public Optional<ModelProvenance> getTribuoProvenance() {
+        try {
+            OnnxModelMetadata metadata = session.getMetadata();
+            Optional<String> value = metadata.getCustomMetadataValue(ONNXExportable.PROVENANCE_METADATA_FIELD);
+            if (value.isPresent()) {
+                Provenance prov = ONNXExportable.SERIALIZER.deserializeAndUnmarshal(value.get());
+
+                if (prov instanceof ModelProvenance) {
+                    return Optional.of((ModelProvenance) prov);
+                } else {
+                    logger.log(Level.WARNING, "Found invalid provenance object, " + prov.toString());
+                    return Optional.empty();
+                }
+            } else {
+                return Optional.empty();
+            }
+        } catch (OrtException e) {
+            logger.log(Level.WARNING,"ORTException when reading session metadata",e);
+            return Optional.empty();
+        } catch (ProvenanceSerializationException e) {
+            logger.log(Level.WARNING, "Failed to parse provenance from value.",e);
+            return Optional.empty();
+        }
+    }
+
     /**
      * Creates an {@code ONNXExternalModel} by loading the model from disk.
-     * @param factory The output factory to use.
-     * @param featureMapping The feature mapping between Tribuo names and ONNX integer ids.
-     * @param outputMapping The output mapping between Tribuo outputs and ONNX integer ids.
+     *
+     * @param factory            The output factory to use.
+     * @param featureMapping     The feature mapping between Tribuo names and ONNX integer ids.
+     * @param outputMapping      The output mapping between Tribuo outputs and ONNX integer ids.
      * @param featureTransformer The transformation function for the features.
-     * @param outputTransformer The transformation function for the outputs.
-     * @param opts The session options for the ONNX model.
-     * @param filename The model path.
-     * @param inputName The name of the input node.
-     * @param <T> The type of the output.
+     * @param outputTransformer  The transformation function for the outputs.
+     * @param opts               The session options for the ONNX model.
+     * @param filename           The model path.
+     * @param inputName          The name of the input node.
+     * @param <T>                The type of the output.
      * @return An ONNXExternalModel ready to score new inputs.
      * @throws OrtException If the onnx-runtime native library call failed.
      */
     public static <T extends Output<T>> ONNXExternalModel<T> createOnnxModel(OutputFactory<T> factory,
                                                                              Map<String, Integer> featureMapping,
-                                                                             Map<T,Integer> outputMapping,
+                                                                             Map<T, Integer> outputMapping,
                                                                              ExampleTransformer featureTransformer,
                                                                              OutputTransformer<T> outputTransformer,
                                                                              OrtSession.SessionOptions opts,
                                                                              String filename, String inputName) throws OrtException {
         Path path = Paths.get(filename);
-        return createOnnxModel(factory,featureMapping,outputMapping,featureTransformer,outputTransformer,
-                opts,path,inputName);
+        return createOnnxModel(factory, featureMapping, outputMapping, featureTransformer, outputTransformer,
+                opts, path, inputName);
     }
 
     /**
      * Creates an {@code ONNXExternalModel} by loading the model from disk.
-     * @param factory The output factory to use.
-     * @param featureMapping The feature mapping between Tribuo names and ONNX integer ids.
-     * @param outputMapping The output mapping between Tribuo outputs and ONNX integer ids.
+     *
+     * @param factory            The output factory to use.
+     * @param featureMapping     The feature mapping between Tribuo names and ONNX integer ids.
+     * @param outputMapping      The output mapping between Tribuo outputs and ONNX integer ids.
      * @param featureTransformer The transformation function for the features.
-     * @param outputTransformer The transformation function for the outputs.
-     * @param opts The session options for the ONNX model.
-     * @param path The model path.
-     * @param inputName The name of the input node.
-     * @param <T> The type of the output.
+     * @param outputTransformer  The transformation function for the outputs.
+     * @param opts               The session options for the ONNX model.
+     * @param path               The model path.
+     * @param inputName          The name of the input node.
+     * @param <T>                The type of the output.
      * @return An ONNXExternalModel ready to score new inputs.
      * @throws OrtException If the onnx-runtime native library call failed.
      */
     public static <T extends Output<T>> ONNXExternalModel<T> createOnnxModel(OutputFactory<T> factory,
                                                                              Map<String, Integer> featureMapping,
-                                                                             Map<T,Integer> outputMapping,
+                                                                             Map<T, Integer> outputMapping,
                                                                              ExampleTransformer featureTransformer,
                                                                              OutputTransformer<T> outputTransformer,
                                                                              OrtSession.SessionOptions opts,
@@ -286,30 +338,30 @@ public static <T extends Output<T>> ONNXExternalModel<T> createOnnxModel(OutputF
             byte[] modelArray = Files.readAllBytes(path);
             URL provenanceLocation = path.toUri().toURL();
             ImmutableFeatureMap featureMap = ExternalModel.createFeatureMap(featureMapping.keySet());
-            ImmutableOutputInfo<T> outputInfo = ExternalModel.createOutputInfo(factory,outputMapping);
+            ImmutableOutputInfo<T> outputInfo = ExternalModel.createOutputInfo(factory, outputMapping);
             OffsetDateTime now = OffsetDateTime.now();
             ExternalTrainerProvenance trainerProvenance = new ExternalTrainerProvenance(provenanceLocation);
-            DatasetProvenance datasetProvenance = new ExternalDatasetProvenance("unknown-external-data",factory,false,featureMapping.size(),outputMapping.size());
+            DatasetProvenance datasetProvenance = new ExternalDatasetProvenance("unknown-external-data", factory, false, featureMapping.size(), outputMapping.size());
             HashMap<String, Provenance> runProvenance = new HashMap<>();
             runProvenance.put("input-name", new StringProvenance("input-name", inputName));
             try (OrtEnvironment env = OrtEnvironment.getEnvironment();
                  OrtSession session = env.createSession(modelArray)) {
                 OnnxModelMetadata metadata = session.getMetadata();
-                runProvenance.put("model-producer", new StringProvenance("model-producer",metadata.getProducerName()));
-                runProvenance.put("model-domain", new StringProvenance("model-domain",metadata.getDomain()));
-                runProvenance.put("model-description", new StringProvenance("model-description",metadata.getDescription()));
-                runProvenance.put("model-graphname", new StringProvenance("model-graphname",metadata.getGraphName()));
-                runProvenance.put("model-version", new LongProvenance("model-version",metadata.getVersion()));
-                for (Map.Entry<String,String> e : metadata.getCustomMetadata().entrySet()) {
-                    String keyName = "model-metadata-"+e.getKey();
-                    runProvenance.put(keyName, new StringProvenance(keyName,e.getValue()));
+                runProvenance.put("model-producer", new StringProvenance("model-producer", metadata.getProducerName()));
+                runProvenance.put("model-domain", new StringProvenance("model-domain", metadata.getDomain()));
+                runProvenance.put("model-description", new StringProvenance("model-description", metadata.getDescription()));
+                runProvenance.put("model-graphname", new StringProvenance("model-graphname", metadata.getGraphName()));
+                runProvenance.put("model-version", new LongProvenance("model-version", metadata.getVersion()));
+                for (Map.Entry<String, String> e : metadata.getCustomMetadata().entrySet()) {
+                    String keyName = "model-metadata-" + e.getKey();
+                    runProvenance.put(keyName, new StringProvenance(keyName, e.getValue()));
                 }
             } catch (OrtException e) {
                 throw new IllegalArgumentException("Failed to load model and read metadata from path " + path, e);
             }
-            ModelProvenance provenance = new ModelProvenance(ONNXExternalModel.class.getName(),now,datasetProvenance,trainerProvenance,runProvenance);
-            return new ONNXExternalModel<>("external-model",provenance,featureMap,outputInfo,
-                    featureMapping,modelArray,opts,inputName,featureTransformer,outputTransformer);
+            ModelProvenance provenance = new ModelProvenance(ONNXExternalModel.class.getName(), now, datasetProvenance, trainerProvenance, runProvenance);
+            return new ONNXExternalModel<>("external-model", provenance, featureMap, outputInfo,
+                    featureMapping, modelArray, opts, inputName, featureTransformer, outputTransformer);
         } catch (IOException e) {
             throw new IllegalArgumentException("Unable to load model from path " + path, e);
         }
@@ -320,7 +372,7 @@ private void readObject(java.io.ObjectInputStream in) throws IOException, ClassN
         try {
             this.env = OrtEnvironment.getEnvironment();
             this.options = new OrtSession.SessionOptions();
-            this.session = env.createSession(modelArray,options);
+            this.session = env.createSession(modelArray, options);
         } catch (OrtException e) {
             throw new IllegalStateException("Could not construct ONNX Runtime session during deserialization.");
         }
diff --git a/pom.xml b/pom.xml
index 719eba40d..6e3c21dff 100644
--- a/pom.xml
+++ b/pom.xml
@@ -42,12 +42,12 @@
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
 
         <!-- MLRG dependencies -->
-        <olcut.version>5.1.6</olcut.version>
+        <olcut.version>5.2.0</olcut.version>
 
         <!-- 3rd party backend dependencies -->
         <liblinear.version>2.43</liblinear.version>
         <libsvm.version>3.25</libsvm.version>
-        <onnxruntime.version>1.7.0</onnxruntime.version>
+        <onnxruntime.version>1.9.0</onnxruntime.version>
         <tensorflow.version>0.3.1</tensorflow.version>
         <xgboost.version>1.4.1</xgboost.version>
 
@@ -55,7 +55,7 @@
         <junit.version>5.7.1</junit.version>
         <opencsv.version>5.4</opencsv.version>
         <commonsmath.version>3.6.1</commonsmath.version>
-        <protobuf.version>3.16.0</protobuf.version>
+        <protobuf.version>3.17.3</protobuf.version>
 
         <!-- Other properties -->
         <!-- Turn off tests which rely on native code -->
@@ -125,6 +125,11 @@
                 <artifactId>olcut-core</artifactId>
                 <version>${olcut.version}</version>
             </dependency>
+            <dependency>
+                <groupId>com.oracle.labs.olcut</groupId>
+                <artifactId>olcut-config-protobuf</artifactId>
+                <version>${olcut.version}</version>
+            </dependency>
             <dependency>
                 <groupId>com.google.protobuf</groupId>
                 <artifactId>protobuf-java</artifactId>

From 7d549574ab2acc49688c4f0add86887075e6ca72 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Thu, 30 Sep 2021 14:27:00 -0400
Subject: [PATCH 2/8] Adding Model.castModel so users don't have an unchecked
 cast in their code.

---
 Core/src/main/java/org/tribuo/Model.java | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/Core/src/main/java/org/tribuo/Model.java b/Core/src/main/java/org/tribuo/Model.java
index ea14efe85..1f42fbcf8 100644
--- a/Core/src/main/java/org/tribuo/Model.java
+++ b/Core/src/main/java/org/tribuo/Model.java
@@ -280,7 +280,7 @@ public Model<T> copy() {
     }
 
     /**
-     * Copies a model, replacing it's provenance and name with the supplied values.
+     * Copies a model, replacing its provenance and name with the supplied values.
      * <p>
      * Used to provide the provenance removal functionality.
      * @param newName The new name.
@@ -297,5 +297,24 @@ public String toString() {
             return provenanceOutput;
         }
     }
+
+    /**
+     * Casts the model to the specified output type, assuming it is valid.
+     * <p>
+     * If it's not valid, throws {@link ClassCastException}.
+     * @param inputModel The model to cast.
+     * @param outputType The output type to cast to.
+     * @param <T> The output type.
+     * @return The model cast to the correct value.
+     */
+    public static <T extends Output<T>> Model<T> castModel(Model<?> inputModel, Class<T> outputType) {
+        if (inputModel.validate(outputType)) {
+            @SuppressWarnings("unchecked") // guarded by validate
+            Model<T> castedModel = (Model<T>) inputModel;
+            return castedModel;
+        } else {
+            throw new ClassCastException("Attempted to cast model to " + outputType.getName() + " which is not valid for model " + inputModel.toString());
+        }
+    }
     
 }

From 172f9b0c49a3f66bf4c6ddaadb5ce0de75c37f15 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 8 Oct 2021 12:34:16 -0400
Subject: [PATCH 3/8] Fixing a bug in IndependentMultiLabelTrainer's
 provenance.

---
 .../baseline/IndependentMultiLabelTrainer.java    | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/MultiLabel/Core/src/main/java/org/tribuo/multilabel/baseline/IndependentMultiLabelTrainer.java b/MultiLabel/Core/src/main/java/org/tribuo/multilabel/baseline/IndependentMultiLabelTrainer.java
index a814cca63..83e12d065 100644
--- a/MultiLabel/Core/src/main/java/org/tribuo/multilabel/baseline/IndependentMultiLabelTrainer.java
+++ b/MultiLabel/Core/src/main/java/org/tribuo/multilabel/baseline/IndependentMultiLabelTrainer.java
@@ -81,8 +81,19 @@ public Model<MultiLabel> train(Dataset<MultiLabel> examples, Map<String, Provena
         }
         ArrayList<Model<Label>> modelsList = new ArrayList<>();
         ArrayList<Label> labelList = new ArrayList<>();
+
+        // Build provenance
         DatasetProvenance datasetProvenance = examples.getProvenance();
-        //TODO supply more suitable provenance showing it's a single dimension out of many.
+        TrainerProvenance trainerProvenance;
+        // Construct the trainer provenance including the inner trainer invocation count field
+        synchronized (innerTrainer) {
+            trainerProvenance = getProvenance();
+            trainInvocationCounter++;
+        }
+        //TODO supply more suitable provenance showing there are multiple models, one per dimension.
+        ModelProvenance provenance = new ModelProvenance(IndependentMultiLabelModel.class.getName(), OffsetDateTime.now(), datasetProvenance, trainerProvenance, runProvenance);
+
+        // Construct binarised training data
         MutableDataset<Label> trainingData = new MutableDataset<>(datasetProvenance, new LabelFactory());
         for (Example<MultiLabel> e : examples) {
             trainingData.add(new BinaryExample(e, MultiLabel.NEGATIVE_LABEL));
@@ -100,8 +111,6 @@ public Model<MultiLabel> train(Dataset<MultiLabel> examples, Map<String, Provena
             trainingData.regenerateOutputInfo();
             modelsList.add(innerTrainer.train(trainingData));
         }
-        ModelProvenance provenance = new ModelProvenance(IndependentMultiLabelModel.class.getName(), OffsetDateTime.now(), datasetProvenance, getProvenance(), runProvenance);
-        trainInvocationCounter++;
         return new IndependentMultiLabelModel(labelList,modelsList,provenance,featureMap,labelInfo);
     }
 

From ba3eb3c30b21238619e41a9197565b4ba6657ae6 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 8 Oct 2021 18:27:08 -0400
Subject: [PATCH 4/8] Adding provenance export to LinearSGDModel and
 SparseLinearModel.

---
 .../sgd/linear/TestSGDLinear.java             | 21 +++++++++++++++----
 .../common/sgd/AbstractLinearSGDModel.java    |  9 ++++++++
 .../multilabel/sgd/linear/TestSGDLinear.java  | 12 +++++++++++
 .../regression/sgd/linear/TestSGDLinear.java  | 12 +++++++++++
 .../regression/slm/SparseLinearModel.java     |  8 +++++++
 .../org/tribuo/regression/slm/TestSLM.java    | 13 ++++++++++++
 6 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
index fa270e5ce..3e80604f8 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java
@@ -46,6 +46,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.test.Helpers;
 
 import java.io.IOException;
@@ -55,11 +56,15 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestSGDLinear {
@@ -82,10 +87,10 @@ public static LinearSGDModel testSGDLinear(Pair<Dataset<Label>,Dataset<Label>> p
         LabelEvaluation evaluation = e.evaluate(m,p.getB());
         Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
         Assertions.assertNotNull(features);
-        Assertions.assertFalse(features.isEmpty());
+        assertFalse(features.isEmpty());
         features = m.getTopFeatures(-1);
         Assertions.assertNotNull(features);
-        Assertions.assertFalse(features.isEmpty());
+        assertFalse(features.isEmpty());
         return m;
     }
 
@@ -152,6 +157,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 }
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);
@@ -209,10 +222,10 @@ public void testSerializedModel(String resourceName) throws IOException, ClassNo
                LabelEvaluation evaluation = e.evaluate(m,LabelledDataGenerator.denseTrainTest().getB());
                Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
                Assertions.assertNotNull(features);
-               Assertions.assertFalse(features.isEmpty());
+               assertFalse(features.isEmpty());
                features = m.getTopFeatures(-1);
                Assertions.assertNotNull(features);
-               Assertions.assertFalse(features.isEmpty());
+               assertFalse(features.isEmpty());
            } else {
                fail("Invalid model type found, expected Label");
            }
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
index bd571c022..2469dff11 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
@@ -31,6 +31,7 @@
 import org.tribuo.math.LinearParameters;
 import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
 import org.tribuo.onnx.ONNXOperators;
 import org.tribuo.provenance.ModelProvenance;
 
@@ -179,6 +180,14 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
         builder.setDocString(toString());
         builder.addOpsetImport(ONNXOperators.getOpsetProto());
         builder.setIrVersion(6);
+
+        // Extract provenance and store in metadata
+        OnnxMl.StringStringEntryProto.Builder metaBuilder = OnnxMl.StringStringEntryProto.newBuilder();
+        metaBuilder.setKey(ONNXExportable.PROVENANCE_METADATA_FIELD);
+        String serializedProvenance = ONNXExportable.SERIALIZER.marshalAndSerialize(getProvenance());
+        metaBuilder.setValue(serializedProvenance);
+        builder.addMetadataProps(metaBuilder.build());
+
         return builder.build();
     }
 
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
index e65568e69..d4aade3a1 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/linear/TestSGDLinear.java
@@ -41,6 +41,7 @@
 import org.tribuo.multilabel.example.MultiLabelDataGenerator;
 import org.tribuo.multilabel.sgd.objectives.Hinge;
 import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.test.Helpers;
 
 import java.io.IOException;
@@ -49,10 +50,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestSGDLinear {
@@ -146,6 +150,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 }
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
index c782c8f66..d99c5378b 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/linear/TestSGDLinear.java
@@ -36,6 +36,7 @@
 import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.regression.RegressionFactory;
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.evaluation.RegressionEvaluation;
@@ -55,12 +56,15 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestSGDLinear {
@@ -134,6 +138,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 assertArrayEquals(tribuo.getOutput().getValues(),external.getOutput().getValues(),1e-5);
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);
diff --git a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
index 23718351a..0a2594a78 100644
--- a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
+++ b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
@@ -237,6 +237,14 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
         builder.setDocString(toString());
         builder.addOpsetImport(ONNXOperators.getOpsetProto());
         builder.setIrVersion(6);
+
+        // Extract provenance and store in metadata
+        OnnxMl.StringStringEntryProto.Builder metaBuilder = OnnxMl.StringStringEntryProto.newBuilder();
+        metaBuilder.setKey(ONNXExportable.PROVENANCE_METADATA_FIELD);
+        String serializedProvenance = ONNXExportable.SERIALIZER.marshalAndSerialize(getProvenance());
+        metaBuilder.setValue(serializedProvenance);
+        builder.addMetadataProps(metaBuilder.build());
+
         return builder.build();
     }
 
diff --git a/Regression/SLM/src/test/java/org/tribuo/regression/slm/TestSLM.java b/Regression/SLM/src/test/java/org/tribuo/regression/slm/TestSLM.java
index 52aa43290..dfc4bd9f4 100644
--- a/Regression/SLM/src/test/java/org/tribuo/regression/slm/TestSLM.java
+++ b/Regression/SLM/src/test/java/org/tribuo/regression/slm/TestSLM.java
@@ -32,6 +32,7 @@
 import org.tribuo.interop.onnx.DenseTransformer;
 import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.interop.onnx.RegressorTransformer;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.regression.RegressionFactory;
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.evaluation.RegressionEvaluation;
@@ -51,11 +52,14 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestSLM {
@@ -130,6 +134,15 @@ public static Model<Regressor> testTrainer(Trainer<Regressor> trainer,
                         assertArrayEquals(tribuo.getOutput().getNames(), external.getOutput().getNames());
                         assertArrayEquals(tribuo.getOutput().getValues(), external.getOutput().getValues(), 1e-5);
                     }
+
+                    // Check that the provenance can be extracted and is the same
+                    ModelProvenance modelProv = m.getProvenance();
+                    Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+                    assertTrue(optProv.isPresent());
+                    ModelProvenance onnxProv = optProv.get();
+                    assertNotSame(onnxProv, modelProv);
+                    assertEquals(modelProv,onnxProv);
+
                     onnxModel.close();
                 } else {
                     logger.warning("ORT based tests only supported on x86_64, found " + arch);

From b6121d7846e62fe2d309e996c8331a0860d0849e Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 12 Oct 2021 18:22:46 -0400
Subject: [PATCH 5/8] Adding provenance export to factorization machines.

---
 .../classification/sgd/fm/TestFMClassification.java  | 12 ++++++++++++
 .../java/org/tribuo/common/sgd/AbstractFMModel.java  |  9 +++++++++
 .../tribuo/multilabel/sgd/fm/TestFMMultiLabel.java   | 12 ++++++++++++
 .../tribuo/regression/sgd/fm/TestFMRegression.java   | 12 ++++++++++++
 4 files changed, 45 insertions(+)

diff --git a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
index 2d0b44dde..4eb0b845d 100644
--- a/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
+++ b/Classification/SGD/src/test/java/org/tribuo/classification/sgd/fm/TestFMClassification.java
@@ -46,6 +46,7 @@
 import org.tribuo.interop.onnx.LabelTransformer;
 import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.test.Helpers;
 
 import java.io.IOException;
@@ -54,12 +55,15 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestFMClassification {
@@ -176,6 +180,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 }
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);
diff --git a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
index df4396422..1ed2933f0 100644
--- a/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
+++ b/Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
@@ -29,6 +29,7 @@
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.math.la.Tensor;
 import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
 import org.tribuo.onnx.ONNXOperators;
 import org.tribuo.provenance.ModelProvenance;
 
@@ -204,6 +205,14 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
         builder.setDocString(toString());
         builder.addOpsetImport(ONNXOperators.getOpsetProto());
         builder.setIrVersion(6);
+
+        // Extract provenance and store in metadata
+        OnnxMl.StringStringEntryProto.Builder metaBuilder = OnnxMl.StringStringEntryProto.newBuilder();
+        metaBuilder.setKey(ONNXExportable.PROVENANCE_METADATA_FIELD);
+        String serializedProvenance = ONNXExportable.SERIALIZER.marshalAndSerialize(getProvenance());
+        metaBuilder.setValue(serializedProvenance);
+        builder.addMetadataProps(metaBuilder.build());
+
         return builder.build();
     }
 
diff --git a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
index dc9638c91..948065bfd 100644
--- a/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
+++ b/MultiLabel/SGD/src/test/java/org/tribuo/multilabel/sgd/fm/TestFMMultiLabel.java
@@ -41,6 +41,7 @@
 import org.tribuo.multilabel.example.MultiLabelDataGenerator;
 import org.tribuo.multilabel.sgd.objectives.BinaryCrossEntropy;
 import org.tribuo.multilabel.sgd.objectives.Hinge;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.test.Helpers;
 
 import java.io.IOException;
@@ -49,10 +50,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestFMMultiLabel {
@@ -146,6 +150,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 }
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);
diff --git a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
index 09e8e8fd1..496e135ce 100644
--- a/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
+++ b/Regression/SGD/src/test/java/org/tribuo/regression/sgd/fm/TestFMRegression.java
@@ -35,6 +35,7 @@
 import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.interop.onnx.RegressorTransformer;
 import org.tribuo.math.optimisers.AdaGrad;
+import org.tribuo.provenance.ModelProvenance;
 import org.tribuo.regression.RegressionFactory;
 import org.tribuo.regression.Regressor;
 import org.tribuo.regression.evaluation.RegressionEvaluation;
@@ -49,12 +50,15 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class TestFMRegression {
     private static final Logger logger = Logger.getLogger(TestFMRegression.class.getName());
@@ -174,6 +178,14 @@ public void testOnnxSerialization() throws IOException, OrtException {
                 assertArrayEquals(tribuo.getOutput().getValues(),external.getOutput().getValues(),1e-4);
             }
 
+            // Check that the provenance can be extracted and is the same
+            ModelProvenance modelProv = model.getProvenance();
+            Optional<ModelProvenance> optProv = onnxModel.getTribuoProvenance();
+            assertTrue(optProv.isPresent());
+            ModelProvenance onnxProv = optProv.get();
+            assertNotSame(onnxProv, modelProv);
+            assertEquals(modelProv,onnxProv);
+
             onnxModel.close();
         } else {
             logger.warning("ORT based tests only supported on x86_64, found " + arch);

From 7fc7d27e87f63a1b22c44467ac8adb7c4dc3ef96 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Wed, 13 Oct 2021 09:51:41 -0400
Subject: [PATCH 6/8] Fixing some test warnings coming from the new ORT
 version.

---
 .../test/java/org/tribuo/interop/onnx/TestOnnxRuntime.java    | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/Interop/ONNX/src/test/java/org/tribuo/interop/onnx/TestOnnxRuntime.java b/Interop/ONNX/src/test/java/org/tribuo/interop/onnx/TestOnnxRuntime.java
index 75fc77f0c..f648e0cc5 100644
--- a/Interop/ONNX/src/test/java/org/tribuo/interop/onnx/TestOnnxRuntime.java
+++ b/Interop/ONNX/src/test/java/org/tribuo/interop/onnx/TestOnnxRuntime.java
@@ -116,6 +116,8 @@ public void testMNIST() throws IOException, OrtException, URISyntaxException {
             ONNXExternalModel<Label> transposedMNISTLR = ONNXExternalModel.createOnnxModel(
                     labelFactory, featureMapping, outputMapping, denseTransformer,
                     labelTransformer, sessionOptions, testResource, "float_input");
+            // This model doesn't have a free batch size parameter on the output
+            transposedMNISTLR.setBatchSize(1);
 
             LabelEvaluation evaluation = labelFactory.getEvaluator().evaluate(transposedMNISTLR, transposedMNIST);
             assertEquals(0.967741, evaluation.accuracy(), 1e-6);
@@ -157,6 +159,8 @@ public void testTransposedMNIST() throws IOException, OrtException, URISyntaxExc
             ONNXExternalModel<Label> transposedMNISTLR = ONNXExternalModel.createOnnxModel(
                     labelFactory, featureMapping, outputMapping, denseTransformer,
                     labelTransformer, sessionOptions, testResource, "float_input");
+            // This model doesn't have a free batch size parameter on the output
+            transposedMNISTLR.setBatchSize(1);
 
             LabelEvaluation evaluation = labelFactory.getEvaluator().evaluate(transposedMNISTLR, transposedMNIST);
             assertEquals(0.967741, evaluation.accuracy(), 1e-6);

From 8c65786e0449aaa2c4926c01db7f2720288a6bcb Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Wed, 13 Oct 2021 09:52:00 -0400
Subject: [PATCH 7/8] Updating third party licenses with the new versions.

---
 THIRD_PARTY_LICENSES.txt | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/THIRD_PARTY_LICENSES.txt b/THIRD_PARTY_LICENSES.txt
index f96c4c963..26fc10250 100644
--- a/THIRD_PARTY_LICENSES.txt
+++ b/THIRD_PARTY_LICENSES.txt
@@ -1,4 +1,4 @@
-OLCUT 5.1.6 - BSD License
+OLCUT 5.2.0 - BSD License
 
 Copyright © 2020, 2021 Oracle and/or its affiliates.
 
@@ -56,9 +56,9 @@ LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-LibSVM 3.24 - BSD License
+LibSVM 3.25 - BSD License
 
-Copyright (c) 2000-2019 Chih-Chung Chang and Chih-Jen Lin
+Copyright (c) 2000-2021 Chih-Chung Chang and Chih-Jen Lin
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
@@ -89,7 +89,7 @@ LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-onnx-runtime 1.7.0 - MIT License
+onnx-runtime 1.9.0 - MIT License
 
 MIT License
 
@@ -1908,7 +1908,7 @@ ONNX v1.9.0 - Apache 2.0
    See the License for the specific language governing permissions and
    limitations under the License.
 
-protobuf-java 3.16.0 - 3-clause BSD
+protobuf-java 3.17.3 - 3-clause BSD
 
 Copyright 2008 Google Inc.  All rights reserved.
 

From beba54dd8c16e388b6664b990e43989ace2f449d Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Mon, 18 Oct 2021 10:47:57 -0400
Subject: [PATCH 8/8] Update
 Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java

Co-authored-by: Jack Sullivan <jack.t.sullivan@oracle.com>
---
 .../main/java/org/tribuo/regression/slm/SparseLinearModel.java | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
index 0a2594a78..ea007eefb 100644
--- a/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
+++ b/Regression/SLM/src/main/java/org/tribuo/regression/slm/SparseLinearModel.java
@@ -241,8 +241,7 @@ public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
         // Extract provenance and store in metadata
         OnnxMl.StringStringEntryProto.Builder metaBuilder = OnnxMl.StringStringEntryProto.newBuilder();
         metaBuilder.setKey(ONNXExportable.PROVENANCE_METADATA_FIELD);
-        String serializedProvenance = ONNXExportable.SERIALIZER.marshalAndSerialize(getProvenance());
-        metaBuilder.setValue(serializedProvenance);
+        metaBuilder.setValue(serializeProvenance(getProvenance()));
         builder.addMetadataProps(metaBuilder.build());
 
         return builder.build();