001/*
002 * Copyright (c) 2021, 2022, Oracle and/or its affiliates. All rights reserved.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.tribuo.classification.sgd.fm;
018
019import com.google.protobuf.Any;
020import com.google.protobuf.InvalidProtocolBufferException;
021import org.tribuo.Example;
022import org.tribuo.ImmutableFeatureMap;
023import org.tribuo.ImmutableOutputInfo;
024import org.tribuo.ONNXExportable;
025import org.tribuo.Prediction;
026import org.tribuo.classification.Label;
027import org.tribuo.classification.sgd.protos.FMClassificationModelProto;
028import org.tribuo.common.sgd.AbstractFMModel;
029import org.tribuo.common.sgd.FMParameters;
030import org.tribuo.impl.ModelDataCarrier;
031import org.tribuo.math.Parameters;
032import org.tribuo.math.la.DenseVector;
033import org.tribuo.math.util.VectorNormalizer;
034import org.tribuo.protos.core.ModelProto;
035import org.tribuo.provenance.ModelProvenance;
036import org.tribuo.util.onnx.ONNXNode;
037
038import java.util.LinkedHashMap;
039import java.util.Map;
040
041/**
042 * The inference time version of a factorization machine trained using SGD.
043 * <p>
044 * See:
045 * <pre>
046 * Rendle, S.
047 * Factorization machines.
048 * 2010 IEEE International Conference on Data Mining
049 * </pre>
050 */
051public class FMClassificationModel extends AbstractFMModel<Label> implements ONNXExportable {
052    private static final long serialVersionUID = 1L;
053
054    /**
055     * Protobuf serialization version.
056     */
057    public static final int CURRENT_VERSION = 0;
058
059    private final VectorNormalizer normalizer;
060
061    /**
062     * Constructs a classification factorization machine trained via SGD.
063     * @param name The model name.
064     * @param provenance The model provenance.
065     * @param featureIDMap The feature domain.
066     * @param outputIDInfo The output domain.
067     * @param parameters The model parameters.
068     * @param normalizer The normalization function.
069     * @param generatesProbabilities Does this model generate probabilities?
070     */
071    FMClassificationModel(String name, ModelProvenance provenance,
072                          ImmutableFeatureMap featureIDMap, ImmutableOutputInfo<Label> outputIDInfo,
073                          FMParameters parameters, VectorNormalizer normalizer, boolean generatesProbabilities) {
074        super(name, provenance, featureIDMap, outputIDInfo, parameters, generatesProbabilities);
075        this.normalizer = normalizer;
076    }
077
078    /**
079     * Deserialization factory.
080     * @param version The serialized object version.
081     * @param className The class name.
082     * @param message The serialized data.
083     * @throws InvalidProtocolBufferException If the protobuf could not be parsed from the {@code message}.
084     * @return The deserialized object.
085     */
086    public static FMClassificationModel deserializeFromProto(int version, String className, Any message) throws InvalidProtocolBufferException {
087        if (version < 0 || version > CURRENT_VERSION) {
088            throw new IllegalArgumentException("Unknown version " + version + ", this class supports at most version " + CURRENT_VERSION);
089        }
090        FMClassificationModelProto proto = message.unpack(FMClassificationModelProto.class);
091
092        ModelDataCarrier<?> carrier = ModelDataCarrier.deserialize(proto.getMetadata());
093        if (!carrier.outputDomain().getOutput(0).getClass().equals(Label.class)) {
094            throw new IllegalStateException("Invalid protobuf, output domain is not a label domain, found " + carrier.outputDomain().getClass());
095        }
096        @SuppressWarnings("unchecked") // guarded by getClass
097        ImmutableOutputInfo<Label> outputDomain = (ImmutableOutputInfo<Label>) carrier.outputDomain();
098
099        Parameters params = Parameters.deserialize(proto.getParams());
100        if (!(params instanceof FMParameters)) {
101            throw new IllegalStateException("Invalid protobuf, parameters must be FMParameters, found " + params.getClass());
102        }
103
104        VectorNormalizer normalizer = VectorNormalizer.deserialize(proto.getNormalizer());
105
106        return new FMClassificationModel(carrier.name(),carrier.provenance(),carrier.featureDomain(),outputDomain,(FMParameters) params, normalizer, carrier.generatesProbabilities());
107    }
108
109    @Override
110    public Prediction<Label> predict(Example<Label> example) {
111        PredAndActive predTuple = predictSingle(example);
112        DenseVector prediction = predTuple.prediction;
113        prediction.normalize(normalizer);
114
115        double maxScore = Double.NEGATIVE_INFINITY;
116        Label maxLabel = null;
117        Map<String,Label> predMap = new LinkedHashMap<>();
118        for (int i = 0; i < prediction.size(); i++) {
119            String labelName = outputIDInfo.getOutput(i).getLabel();
120            double score = prediction.get(i);
121            Label label = new Label(labelName, score);
122            predMap.put(labelName,label);
123            if (score > maxScore) {
124                maxScore = score;
125                maxLabel = label;
126            }
127        }
128        return new Prediction<>(maxLabel, predMap, predTuple.numActiveFeatures, example, generatesProbabilities);
129    }
130
131    @Override
132    public ModelProto serialize() {
133        ModelDataCarrier<Label> carrier = createDataCarrier();
134        FMClassificationModelProto.Builder modelBuilder = FMClassificationModelProto.newBuilder();
135        modelBuilder.setMetadata(carrier.serialize());
136        modelBuilder.setParams(modelParameters.serialize());
137        modelBuilder.setNormalizer(normalizer.serialize());
138
139        ModelProto.Builder builder = ModelProto.newBuilder();
140        builder.setVersion(CURRENT_VERSION);
141        builder.setClassName(FMClassificationModel.class.getName());
142        builder.setSerializedData(Any.pack(modelBuilder.build()));
143
144        return builder.build();
145    }
146
147    @Override
148    protected FMClassificationModel copy(String newName, ModelProvenance newProvenance) {
149        return new FMClassificationModel(newName,newProvenance,featureIDMap,outputIDInfo,(FMParameters)modelParameters.copy(),normalizer,generatesProbabilities);
150    }
151
152    @Override
153    protected String getDimensionName(int index) {
154        return outputIDInfo.getOutput(index).getLabel();
155    }
156
157    @Override
158    protected String onnxModelName() {
159        return "FMClassificationModel";
160    }
161
162    @Override
163    protected ONNXNode onnxOutput(ONNXNode input) {
164        return normalizer.exportNormalizer(input);
165    }
166
167}