xapian-letor/bin/xapian-train.cc

   1 /** @file xapian-train.cc
   2  * @brief Command line tool to train and save the LTR model
   3  */
   4 /* Copyright (C) 2004,2005,2006,2007,2008,2009,2010,2015 Olly Betts
   5  * Copyright (C) 2011 Parth Gupta
   6  * Copyright (C) 2016 Ayush Tomar
   7  *
   8  * This program is free software; you can redistribute it and/or
   9  * modify it under the terms of the GNU General Public License as
  10  * published by the Free Software Foundation; either version 2 of the
  11  * License, or (at your option) any later version.
  12  *
  13  * This program is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16  * GNU General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License
  19  * along with this program; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301
  21  * USA
  22  */
  23
  24 #include <config.h>
  25
  26 #include <xapian.h>
  27 #include <xapian-letor.h>
  28
  29 #include <iostream>
  30 #include <string>
  31
  32 #include "gnu_getopt.h"
  33
  34 using namespace std;
  35
  36 #define PROG_NAME "xapian-train"
  37 #define PROG_DESC "Command line tool to train and save the LTR model"
  38
  39 #define OPT_HELP 1
  40 #define OPT_VERSION 2
  41
  42 static void show_usage() {
  43     cout << "Usage: " PROG_NAME " [OPTIONS] PATH_TO_TRAINING_FILE MODEL_METADATA_KEY\n"
  44     "Options:\n"
  45     "  -d, --db=DIRECTORY  path to database to search\n"
  46     "      --help          display this help and exit\n"
  47     "      --version       output version information and exit\n";
  48 }
  49
  50 int
  51 main(int argc, char **argv)
  52 try {
  53     const char * opts = "d:h:v";
  54     static const struct option long_opts[] = {
  55         { "db",         required_argument, 0, 'd' },
  56         { "help",       no_argument, 0, OPT_HELP },
  57         { "version",    no_argument, 0, OPT_VERSION },
  58         { NULL,         0, 0, 0}
  59     };
  60
  61     bool have_database = false;
  62
  63     string db_path;
  64
  65     int c;
  66     while ((c = gnu_getopt_long(argc, argv, opts, long_opts, 0)) != -1) {
  67         switch (c) {
  68             case 'd':
  69                 db_path = optarg;
  70                 have_database = true;
  71                 break;
  72             case OPT_HELP:
  73                 cout << PROG_NAME " - " PROG_DESC "\n\n";
  74                 show_usage();
  75                 exit(0);
  76             case OPT_VERSION:
  77                 cout << PROG_NAME " - " PACKAGE_STRING << endl;
  78                 exit(0);
  79             case ':': // missing parameter
  80             case '?': // unknown option
  81                 show_usage();
  82                 exit(1);
  83         }
  84     }
  85
  86     if (argc - optind != 2) {
  87         show_usage();
  88         exit(1);
  89     }
  90
  91     string trainingfile = argv[optind];
  92     string model_metadata_key = argv[optind + 1];
  93
  94     if (!have_database) {
  95         cout << "No database specified so not running the query." << endl;
  96         exit(0);
  97     }
  98
  99     // Initialise Ranker object.
 100     // See Ranker documentation for available Ranker subclass options.
 101     Xapian::Ranker * ranker = new Xapian::ListNETRanker();
 102
 103     // Set database
 104     ranker->set_database_path(db_path);
 105
 106     // Perform training and save model as database metadata with key "model_metadata_key"
 107     ranker->train_model(trainingfile, model_metadata_key);
 108
 109     delete ranker;
 110
 111     cout << flush;
 112
 113 } catch (const Xapian::Error & err) {
 114     cout << err.get_description() << endl;
 115     exit(1);
 116 }