xapian-applications/omega/generate-qrel-file.cc

   1 /** @file
   2  * @brief generates qrel file needed to prepare training file for letor
   3  */
   4 /* Copyright (C) 2017 Vivek Pal
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 2 of the
   9  * License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14  * GNU General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, write to the Free Software Foundation,
  18  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 #include <config.h>
  22
  23 #include <iostream>
  24 #include <fstream>
  25 #include <string>
  26 #include <utility>
  27 #include <vector>
  28
  29 #include "gnu_getopt.h"
  30
  31 #include "clickmodel/simplifieddbn.h"
  32
  33 using namespace std;
  34
  35 #define PROG_NAME "generate-qrel-file"
  36 #define PROG_DESC "Generate qrel file needed to prepare training file for letor"
  37
  38 #define OPT_HELP 1
  39 #define OPT_VERSION 2
  40
  41 static void show_usage() {
  42     cout << "Usage: " PROG_NAME " [OPTIONS] FINAL_LOG QREL_FILE\n\n"
  43 "FINAL_LOG is the path to log file from the 'postprocess' script.\n\n"
  44 "QREL_FILE is the path to save the qrel file to.\n\n"
  45 "Options:\n"
  46 "     --help        display this help and exit\n"
  47 "  --version        output version information and exit" << endl;
  48 }
  49
  50 int
  51 main(int argc, char **argv)
  52 {
  53     const char * opts = "";
  54     static const struct option long_opts[] = {
  55         { "help",       no_argument, 0, OPT_HELP },
  56         { "version",    no_argument, 0, OPT_VERSION },
  57         { NULL,         0, 0, 0}
  58     };
  59
  60     int c;
  61     while ((c = gnu_getopt_long(argc, argv, opts, long_opts, 0)) != -1) {
  62         switch (c) {
  63             case OPT_HELP:
  64                 cout << PROG_NAME " - " PROG_DESC "\n\n";
  65                 show_usage();
  66                 exit(0);
  67             case OPT_VERSION:
  68                 cout << PROG_NAME " - " PACKAGE_STRING << endl;
  69                 exit(0);
  70             default:
  71                 show_usage();
  72                 exit(1);
  73         }
  74     }
  75
  76     if (argc - optind != 2) {
  77         show_usage();
  78         exit(1);
  79     }
  80
  81     string final_log_file = argv[optind];
  82     string qrel_file = argv[optind + 1];
  83
  84     SimplifiedDBN sdbn;
  85
  86     vector<Session> sessions;
  87     try {
  88         sessions = sdbn.build_sessions(final_log_file);
  89     } catch (std::exception &ex) {
  90         cerr << ex.what() << endl;
  91         exit(1);
  92     }
  93
  94     ofstream file_q;
  95     file_q.open(qrel_file, ios::out);
  96
  97     sdbn.train(sessions);
  98
  99     // Extract doc relevances and doc ids from each session and write
 100     // to the qrel file in the required format.
 101     for (auto&& session : sessions) {
 102         vector<pair<string, double>> docid_relevances =
 103             sdbn.get_predicted_relevances(session);
 104
 105         auto reliter = docid_relevances.begin();
 106
 107         for (; reliter != docid_relevances.end(); ++reliter)
 108             file_q << session.get_qid() << " Q0 " << (*reliter).first << ' '
 109                    << (*reliter).second << endl;
 110     }
 111
 112     file_q.close();
 113 }