1 files changed, 38 insertions, 21 deletions

diff --git a/main.cpp b/main.cpp
index 434437f..5c977d2 100644
--- a/main.cpp
+++ b/main.cpp
@@ -10,6 +10,7 @@
 INITIALIZE_EASYLOGGINGPP
 
 using namespace Eigen;
+using namespace std;
 namespace po = boost::program_options;
 
 po::variables_map vm;
@@ -18,7 +19,7 @@ typedef int (*mainFunc)(DataProvider &dp);
 
 int train(DataProvider &dp) {
     RSVM *rsvm;
-    rsvm = RSVM::loadModel(vm["model"].as<std::string>());
+    rsvm = RSVM::loadModel(vm["model"].as<string>());
 
     dp.open();
     DataList D;
@@ -27,31 +28,39 @@ int train(DataProvider &dp) {
     dp.getAllData(D);
     LOG(INFO)<<"Read "<<D.getSize()<<" entries with "<< D.getfSize()<<" features";
     rsvm->train(D);
-    std::vector<double> L;
+    vector<double> L;
     rsvm->predict(D,L);
 
-    rank_accu(D,L);
-
     LOG(INFO)<<"Training finished,saving model";
 
     dp.close();
-    rsvm->saveModel(vm["output"].as<std::string>().c_str());
+    rsvm->saveModel(vm["output"].as<string>().c_str());
     delete rsvm;
     return 0;
 }
 
 int predict(DataProvider &dp) {
     RSVM *rsvm;
-    rsvm = RSVM::loadModel(vm["model"].as<std::string>().c_str());
+    rsvm = RSVM::loadModel(vm["model"].as<string>().c_str());
 
     dp.open();
     DataList D;
-    std::vector<double> L;
+    vector<double> L;
+    CMC cmc;
     LOG(INFO)<<"Prediction started";
 
-    std::ofstream fout;
+    ofstream fout;
     if (vm.count("output"))
-        fout.open(vm["output"].as<std::string>().c_str());
+        fout.open(vm["output"].as<string>().c_str());
+
+    ostream* ot;
+
+    if (vm.count("output")) {
+        fout.open(vm["output"].as<string>().c_str());
+        ot=&fout;
+    }
+    else
+        ot=&cout;
 
     while (!dp.EOFile())
     {
@@ -62,17 +71,24 @@ int predict(DataProvider &dp) {
         if (vm.count("validate"))
         {
             rank_accu(D,L);
+            if (vm.count("cmc"))
+                rank_CMC(D,L,cmc);
         }
 
-        if (vm.count("output"))
-            for (int i=0; i<L.size();++i)
-                fout<<L[i]<<std::endl;
-        else if (!vm.count("validate"))
+        if (vm.count("output") || !vm.count("validate"))
             for (int i=0; i<L.size();++i)
-                std::cout<<L[i]<<std::endl;
+                *ot<<L[i]<<endl;
     }
 
     LOG(INFO)<<"Finished";
+    if (vm.count("cmc"))
+    {
+        LOG(INFO)<< "CMC accounted over " <<cmc.getCount() << " queries";
+        *ot << "CMC"<<endl;
+        vector<double> cur = cmc.getAcc();
+        for (int i = 0;i<CMC_MAX;++i)
+        *ot << cur[i]<<endl;
+    }
     if (vm.count("output"))
         fout.close();
     dp.close();
@@ -88,9 +104,10 @@ int main(int argc, char **argv) {
             ("train,T", "training model")
             ("validate,V", "validate model")
             ("predict,P", "use model for prediction")
-            ("model,m", po::value<std::string>(), "set input model file")
-            ("output,o", po::value<std::string>(), "set output model/prediction file")
-            ("feature,i", po::value<std::string>(), "set input feature file");
+            ("cmc,C", "enable cmc auditing")
+            ("model,m", po::value<string>(), "set input model file")
+            ("output,o", po::value<string>(), "set output model/prediction file")
+            ("feature,i", po::value<string>(), "set input feature file");
 
     // Parsing program options
     po::store(po::parse_command_line(argc, argv, desc), vm);
@@ -98,7 +115,7 @@ int main(int argc, char **argv) {
 
     // Print help if necessary
     if (vm.count("help") || !(vm.count("train") || vm.count("validate") || vm.count("predict"))) {
-        std::cout << desc;
+        cout << desc;
         return 0;
     }
     mainFunc mainf;
@@ -110,10 +127,10 @@ int main(int argc, char **argv) {
     }
     else return 0;
     DataProvider* dp;
-    if (vm["feature"].as<std::string>().find(".rid") == std::string::npos)
-        dp = new FileDP(vm["feature"].as<std::string>());
+    if (vm["feature"].as<string>().find(".rid") == string::npos)
+        dp = new FileDP(vm["feature"].as<string>());
     else
-        dp = new RidFileDP(vm["feature"].as<std::string>());
+        dp = new RidFileDP(vm["feature"].as<string>());
     mainf(*dp);
     delete dp;
     return 0;