1 files changed, 21 insertions, 19 deletions

diff --git a/main.cpp b/main.cpp
index d5c6f69..cf1dd14 100644
--- a/main.cpp
+++ b/main.cpp
@@ -24,7 +24,7 @@ int train(DataProvider &dp) {
     DataList D;
 
     LOG(INFO)<<"Training started";
-    dp.getDataSet(D);
+    dp.getAllData(D);
     LOG(INFO)<<"Read "<<D.getSize()<<" entries with "<< D.getfSize()<<" features";
     rsvm->train(D);
     std::vector<double> L;
@@ -49,30 +49,32 @@ int predict(DataProvider &dp) {
     std::vector<double> L;
     LOG(INFO)<<"Prediction started";
 
-    dp.getDataSet(D);
-    LOG(INFO)<<"Read "<<D.getSize()<<" entries with "<< D.getfSize()<<" features";
-    rsvm->predict(D,L);
+    std::ofstream fout;
+    if (vm.count("output"))
+        fout.open(vm["output"].as<std::string>().c_str());
 
-    if (vm.count("validate"))
+    while (!dp.EOFile())
     {
-        rank_accu(D,L);
+        dp.getDataSet(D);
+        LOG(INFO)<<"Read "<<D.getSize()<<" entries with "<< D.getfSize()<<" features";
+        rsvm->predict(D,L);
+
+        if (vm.count("validate"))
+        {
+            rank_accu(D,L);
+        }
+
+        if (vm.count("output"))
+            for (int i=0; i<L.size();++i)
+                fout<<L[i]<<std::endl;
+        else if (!vm.count("validate"))
+            for (int i=0; i<L.size();++i)
+                std::cout<<L[i]<<std::endl;
     }
 
+    LOG(INFO)<<"Finished";
     if (vm.count("output"))
-    {
-        LOG(INFO)<<"Finished,saving prediction";
-        std::ofstream fout(vm["output"].as<std::string>().c_str());
-
-        for (int i=0; i<L.size();++i)
-            fout<<L[i]<<std::endl;
         fout.close();
-    }
-    else if (!vm.count("validate"))
-    {
-        LOG(INFO)<<"Finished";
-        for (int i=0; i<L.size();++i)
-            std::cout<<L[i]<<std::endl;
-    }
     dp.close();
     delete rsvm;
     return 0;