001/*
002 * Copyright (c) 2015, 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.xgboost;
018
019import com.google.protobuf.Any;
020import com.google.protobuf.ByteString;
021import org.tribuo.Example;
022import org.tribuo.ImmutableOutputInfo;
023import org.tribuo.Prediction;
024import org.tribuo.classification.Label;
025import org.tribuo.common.xgboost.XGBoostOutputConverter;
026import org.tribuo.common.xgboost.protos.XGBoostOutputConverterProto;
027import org.tribuo.protos.ProtoSerializableClass;
028import org.tribuo.protos.ProtoUtil;
029
030import java.util.ArrayList;
031import java.util.LinkedHashMap;
032import java.util.List;
033
034/**
035 * Converts XGBoost outputs into {@link Label} {@link Prediction}s.
036 */
037@ProtoSerializableClass(version = XGBoostClassificationConverter.CURRENT_VERSION)
038public final class XGBoostClassificationConverter implements XGBoostOutputConverter<Label> {
039    private static final long serialVersionUID = 1L;
040
041    /**
042     * Protobuf serialization version.
043     */
044    public static final int CURRENT_VERSION = 0;
045
046    /**
047     * Constructs an XGBoostClassificationConverter.
048     */
049    public XGBoostClassificationConverter() {}
050
051    /**
052     * Deserialization factory.
053     * @param version The serialized object version.
054     * @param className The class name.
055     * @param message The serialized data.
056     * @return The deserialized object.
057     */
058    public static XGBoostClassificationConverter deserializeFromProto(int version, String className, Any message) {
059        if (version < 0 || version > CURRENT_VERSION) {
060            throw new IllegalArgumentException("Unknown version " + version + ", this class supports at most version " + CURRENT_VERSION);
061        }
062        if (message.getValue() != ByteString.EMPTY) {
063            throw new IllegalArgumentException("Invalid proto");
064        }
065        return new XGBoostClassificationConverter();
066    }
067
068    @Override
069    public boolean generatesProbabilities() {
070        return true;
071    }
072
073    @Override
074    public Prediction<Label> convertOutput(ImmutableOutputInfo<Label> info, List<float[]> probabilitiesList, int numValidFeatures, Example<Label> example) {
075        if (probabilitiesList.size() != 1) {
076            throw new IllegalArgumentException("XGBoostClassificationConverter only expects a single model output.");
077        }
078        double maxScore = Double.NEGATIVE_INFINITY;
079        Label maxLabel = null;
080        LinkedHashMap<String,Label> probMap = new LinkedHashMap<>();
081        float[] probabilities = probabilitiesList.get(0);
082
083        for (int i = 0; i < probabilities.length; i++) {
084            String name = info.getOutput(i).getLabel();
085            Label label = new Label(name, probabilities[i]);
086            probMap.put(name, label);
087            if (label.getScore() > maxScore) {
088                maxScore = label.getScore();
089                maxLabel = label;
090            }
091        }
092
093        return new Prediction<>(maxLabel,probMap,numValidFeatures,example,true);
094    }
095
096    @Override
097    public List<Prediction<Label>> convertBatchOutput(ImmutableOutputInfo<Label> info, List<float[][]> probabilitiesList, int[] numValidFeatures, Example<Label>[] examples) {
098        if (probabilitiesList.size() != 1) {
099            throw new IllegalArgumentException("XGBoostClassificationConverter only expects a single model output.");
100        }
101        float[][] probabilities = probabilitiesList.get(0);
102
103        List<Prediction<Label>> predictions = new ArrayList<>();
104        for (int i = 0; i < probabilities.length; i++) {
105            double maxScore = Double.NEGATIVE_INFINITY;
106            Label maxLabel = null;
107            LinkedHashMap<String, Label> probMap = new LinkedHashMap<>();
108            for (int j = 0; j < probabilities[i].length; j++) {
109                String name = info.getOutput(j).getLabel();
110                Label label = new Label(name, probabilities[i][j]);
111                probMap.put(name, label);
112                if (label.getScore() > maxScore) {
113                    maxScore = label.getScore();
114                    maxLabel = label;
115                }
116            }
117
118            predictions.add(new Prediction<>(maxLabel, probMap, numValidFeatures[i], examples[i], true));
119        }
120
121        return predictions;
122    }
123
124    @Override
125    public XGBoostOutputConverterProto serialize() {
126        return ProtoUtil.serialize(this);
127    }
128
129    @Override
130    public Class<Label> getTypeWitness() {
131        return Label.class;
132    }
133}