remove cerr/cout from fasttext test

Summary: See title. Reviewed By: EdouardGrave Differential Revision: D6621580 fbshipit-source-id: 271e9c3c2bc01bb6904c6b7ef829c2574569a6d9
kakonyitibor · Dec 22, 2017 · 9210895 · 9210895
1 parent fce60af
commit 9210895
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 10 deletions.
diff --git a/src/fasttext.cc b/src/fasttext.cc
@@ -360,7 +360,9 @@ void FastText::skipgram(Model& model, real lr,
   }
 }
 
-void FastText::test(std::istream& in, int32_t k) {
+std::tuple<int64_t, double, double> FastText::test(
+    std::istream& in,
+    int32_t k) {
   int32_t nexamples = 0, nlabels = 0;
   double precision = 0.0;
   std::vector<int32_t> line, labels;
@@ -379,11 +381,8 @@ void FastText::test(std::istream& in, int32_t k) {
       nlabels += labels.size();
     }
   }
-  std::cout << "N" << "\t" << nexamples << std::endl;
-  std::cout << std::setprecision(3);
-  std::cout << "P@" << k << "\t" << precision / (k * nexamples) << std::endl;
-  std::cout << "R@" << k << "\t" << precision / nlabels << std::endl;
-  std::cerr << "Number of examples: " << nexamples << std::endl;
+  return std::tuple<int64_t, double, double>(
+      nexamples, precision / (k * nexamples), precision / nlabels);
 }
 
 void FastText::predict(std::istream& in, int32_t k,

diff --git a/src/fasttext.h b/src/fasttext.h
@@ -16,6 +16,8 @@
 #include <set>
 #include <chrono>
 #include <iostream>
+#include <queue>
+#include <tuple>
 
 #include "args.h"
 #include "dictionary.h"
@@ -91,7 +93,7 @@ class FastText {
   std::vector<int32_t> selectEmbeddings(int32_t) const;
   void getSentenceVector(std::istream&, Vector&);
   void quantize(const Args);
-  void test(std::istream&, int32_t);
+  std::tuple<int64_t, double, double> test(std::istream&, int32_t);
   void predict(std::istream&, int32_t, bool);
   void predict(
       std::istream&,

diff --git a/src/main.cc b/src/main.cc
@@ -8,6 +8,8 @@
  */
 
 #include <iostream>
+#include <queue>
+#include <iomanip>
 #include "fasttext.h"
 #include "args.h"
 
@@ -132,19 +134,24 @@ void test(const std::vector<std::string>& args) {
   FastText fasttext;
   fasttext.loadModel(args[2]);
 
+  std::tuple<int64_t, double, double> result;
   std::string infile = args[3];
   if (infile == "-") {
-    fasttext.test(std::cin, k);
+    result = fasttext.test(std::cin, k);
   } else {
     std::ifstream ifs(infile);
     if (!ifs.is_open()) {
       std::cerr << "Test file cannot be opened!" << std::endl;
       exit(EXIT_FAILURE);
     }
-    fasttext.test(ifs, k);
+    result = fasttext.test(ifs, k);
     ifs.close();
   }
-  exit(0);
+  std::cout << "N" << "\t" << std::get<0>(result) << std::endl;
+  std::cout << std::setprecision(3);
+  std::cout << "P@" << k << "\t" << std::get<1>(result) << std::endl;
+  std::cout << "R@" << k << "\t" << std::get<2>(result) << std::endl;
+  std::cerr << "Number of examples: " << std::get<0>(result) << std::endl;
 }
 
 void predict(const std::vector<std::string>& args) {