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.sgd.crf;
018
019import com.oracle.labs.mlrg.olcut.config.ConfigurationManager;
020import com.oracle.labs.mlrg.olcut.config.Option;
021import com.oracle.labs.mlrg.olcut.config.Options;
022import com.oracle.labs.mlrg.olcut.config.UsageException;
023import com.oracle.labs.mlrg.olcut.util.LabsLogFormatter;
024import org.tribuo.classification.Label;
025import org.tribuo.classification.sequence.LabelSequenceEvaluation;
026import org.tribuo.classification.sequence.LabelSequenceEvaluator;
027import org.tribuo.classification.sequence.example.SequenceDataGenerator;
028import org.tribuo.hash.HashCodeHasher;
029import org.tribuo.hash.HashingOptions.ModelHashingType;
030import org.tribuo.hash.MessageDigestHasher;
031import org.tribuo.math.StochasticGradientOptimiser;
032import org.tribuo.math.optimisers.GradientOptimiserOptions;
033import org.tribuo.sequence.HashingSequenceTrainer;
034import org.tribuo.sequence.SequenceDataset;
035import org.tribuo.sequence.SequenceTrainer;
036import org.tribuo.util.Util;
037
038import java.io.BufferedInputStream;
039import java.io.IOException;
040import java.io.ObjectInputStream;
041import java.io.ObjectOutputStream;
042import java.nio.file.Files;
043import java.nio.file.Path;
044import java.util.logging.Logger;
045
046/**
047 * Build and run a sequence classifier on a generated dataset.
048 */
049public class SeqTest {
050
051    private static final Logger logger = Logger.getLogger(SeqTest.class.getName());
052
053    /**
054     * Command line options.
055     */
056    public static class CRFOptions implements Options {
057        @Override
058        public String getOptionsDescription() {
059            return "Tests a linear chain CRF model on the specified dataset.";
060        }
061
062        /**
063         * Gradient descent options.
064         */
065        public GradientOptimiserOptions gradientOptions;
066        /**
067         * Name of the example dataset, options are {gorilla}.
068         */
069        @Option(charName = 'd', longName = "dataset-name", usage = "Name of the example dataset, options are {gorilla}.")
070        public String datasetName = "";
071        /**
072         * Path to serialize model to.
073         */
074        @Option(charName = 'f', longName = "output-path", usage = "Path to serialize model to.")
075        public Path outputPath;
076        /**
077         * Number of SGD epochs.
078         */
079        @Option(charName = 'i', longName = "epochs", usage = "Number of SGD epochs.")
080        public int epochs = 5;
081        /**
082         * Print out feature, label and other model details.
083         */
084        @Option(charName = 'o', longName = "print-model", usage = "Print out feature, label and other model details.")
085        public boolean logModel = false;
086        /**
087         * Log the objective after n examples.
088         */
089        @Option(charName = 'p', longName = "logging-interval", usage = "Log the objective after <int> examples.")
090        public int loggingInterval = 100;
091        /**
092         * RNG seed.
093         */
094        @Option(charName = 'r', longName = "seed", usage = "RNG seed.")
095        public long seed = 1;
096        /**
097         * Shuffle the data each epoch (default: true).
098         */
099        @Option(longName = "shuffle", usage = "Shuffle the data each epoch (default: true).")
100        public boolean shuffle = true;
101        /**
102         * Path to a serialised SequenceDataset used for training.
103         */
104        @Option(charName = 'u', longName = "train-dataset", usage = "Path to a serialised SequenceDataset used for training.")
105        public Path trainDataset = null;
106        /**
107         * Path to a serialised SequenceDataset used for testing.
108         */
109        @Option(charName = 'v', longName = "test-dataset", usage = "Path to a serialised SequenceDataset used for testing.")
110        public Path testDataset = null;
111        /**
112         * Hash the model during training. Defaults to no hashing.
113         */
114        @Option(longName = "model-hashing-algorithm", usage = "Hash the model during training. Defaults to no hashing.")
115        public ModelHashingType modelHashingAlgorithm = ModelHashingType.NONE;
116        /**
117         * Salt for hashing the model.
118         */
119        @Option(longName = "model-hashing-salt", usage = "Salt for hashing the model.")
120        public String modelHashingSalt = "";
121        /**
122         * Load in the model in protobuf format.
123         */
124        @Option(longName = "protobuf-model", usage = "Load the model from a protobuf. Optional")
125        public boolean protobufFormat;
126    }
127
128    /**
129     * @param args the command line arguments
130     * @throws ClassNotFoundException if it failed to load the model.
131     * @throws IOException            if there is any error reading the examples.
132     */
133    @SuppressWarnings("unchecked") // deserialising a generic dataset.
134    public static void main(String[] args) throws ClassNotFoundException, IOException {
135
136        //
137        // Use the labs format logging.
138        LabsLogFormatter.setAllLogFormatters();
139
140        CRFOptions o = new CRFOptions();
141        ConfigurationManager cm;
142        try {
143            cm = new ConfigurationManager(args, o);
144        } catch (UsageException e) {
145            logger.info(e.getMessage());
146            return;
147        }
148
149        logger.info("Configuring gradient optimiser");
150        StochasticGradientOptimiser grad = o.gradientOptions.getOptimiser();
151
152        logger.info(String.format("Set logging interval to %d", o.loggingInterval));
153
154        SequenceDataset<Label> train;
155        SequenceDataset<Label> test;
156        switch (o.datasetName) {
157            case "Gorilla":
158            case "gorilla":
159                logger.info("Generating gorilla dataset");
160                train = SequenceDataGenerator.generateGorillaDataset(1);
161                test = SequenceDataGenerator.generateGorillaDataset(1);
162                break;
163            default:
164                if ((o.trainDataset != null) && (o.testDataset != null)) {
165                    if (o.protobufFormat) {
166                        logger.info("Loading protobuf format training data from " + o.trainDataset);
167                        SequenceDataset<?> tmpTrain = SequenceDataset.deserializeFromFile(o.trainDataset);
168                        train = SequenceDataset.castDataset(tmpTrain, Label.class);
169                        logger.info(String.format("Loaded %d training examples for %s", train.size(), train.getOutputs().toString()));
170                        logger.info("Found " + train.getFeatureIDMap().size() + " features");
171                        logger.info("Loading protobuf format testing data from " + o.testDataset);
172                        SequenceDataset<?> tmpTest = SequenceDataset.deserializeFromFile(o.testDataset);
173                        test = SequenceDataset.castDataset(tmpTest, Label.class);
174                        logger.info(String.format("Loaded %d testing examples", test.size()));
175                    } else {
176                        logger.info("Loading training data from " + o.trainDataset);
177                        try (ObjectInputStream ois = new ObjectInputStream(new BufferedInputStream(Files.newInputStream(o.trainDataset)));
178                             ObjectInputStream oits = new ObjectInputStream(new BufferedInputStream(Files.newInputStream(o.testDataset)))) {
179                            @SuppressWarnings("unchecked") // deserialising a generic dataset.
180                            SequenceDataset<Label> tmpTrain = (SequenceDataset<Label>) ois.readObject();
181                            train = tmpTrain;
182                            logger.info(String.format("Loaded %d training examples for %s", train.size(), train.getOutputs().toString()));
183                            logger.info("Found " + train.getFeatureIDMap().size() + " features");
184                            logger.info("Loading testing data from " + o.testDataset);
185                            @SuppressWarnings("unchecked") // deserialising a generic dataset.
186                            SequenceDataset<Label> tmpTest = (SequenceDataset<Label>) oits.readObject();
187                            test = tmpTest;
188                            logger.info(String.format("Loaded %d testing examples", test.size()));
189                        }
190                    }
191                } else {
192                    logger.warning("Unknown dataset " + o.datasetName);
193                    logger.info(cm.usage());
194                    return;
195                }
196        }
197
198        SequenceTrainer<Label> trainer = new CRFTrainer(grad, o.epochs, o.loggingInterval, o.seed);
199        ((CRFTrainer) trainer).setShuffle(o.shuffle);
200        switch (o.modelHashingAlgorithm) {
201            case NONE:
202                break;
203            case HC:
204                trainer = new HashingSequenceTrainer<>(trainer, new HashCodeHasher(o.modelHashingSalt));
205                break;
206            case SHA1:
207                trainer = new HashingSequenceTrainer<>(trainer, new MessageDigestHasher("SHA1", o.modelHashingSalt));
208                break;
209            case SHA256:
210                trainer = new HashingSequenceTrainer<>(trainer, new MessageDigestHasher("SHA-256", o.modelHashingSalt));
211                break;
212            default:
213                logger.info("Unknown hasher " + o.modelHashingAlgorithm);
214        }
215
216        logger.info("Training using " + trainer.toString());
217        final long trainStart = System.currentTimeMillis();
218        CRFModel model = (CRFModel) trainer.train(train);
219        final long trainStop = System.currentTimeMillis();
220        logger.info("Finished training classifier " + Util.formatDuration(trainStart, trainStop));
221
222        if (o.logModel) {
223            System.out.println("FeatureMap = " + model.getFeatureIDMap().toString());
224            System.out.println("LabelMap = " + model.getOutputIDInfo().toString());
225            System.out.println("Features - " + model.generateWeightsString());
226        }
227
228        LabelSequenceEvaluator labelEvaluator = new LabelSequenceEvaluator();
229        final long testStart = System.currentTimeMillis();
230        LabelSequenceEvaluation evaluation = labelEvaluator.evaluate(model,test);
231        final long testStop = System.currentTimeMillis();
232        logger.info("Finished evaluating model " + Util.formatDuration(testStart, testStop));
233        System.out.println(evaluation.toString());
234        System.out.println();
235        System.out.println(evaluation.getConfusionMatrix().toString());
236
237        if (o.outputPath != null) {
238            if (o.protobufFormat) {
239                model.serializeToFile(o.outputPath);
240            } else {
241                try (ObjectOutputStream oos = new ObjectOutputStream(Files.newOutputStream(o.outputPath))) {
242                    oos.writeObject(model);
243                }
244            }
245            logger.info("Serialized model to file: " + o.outputPath);
246        }
247    }
248}