1 files changed, 14 insertions, 8 deletions
diff --git a/klm/lm/builder/interpolate.hh b/klm/lm/builder/interpolate.hh
index 9268d404..0acece92 100644
--- a/klm/lm/builder/interpolate.hh
+++ b/klm/lm/builder/interpolate.hh
@@ -1,9 +1,11 @@
-#ifndef LM_BUILDER_INTERPOLATE__
-#define LM_BUILDER_INTERPOLATE__
+#ifndef LM_BUILDER_INTERPOLATE_H
+#define LM_BUILDER_INTERPOLATE_H
 
-#include <stdint.h>
+#include "util/stream/multi_stream.hh"
+
+#include <vector>
 
-#include "lm/builder/multi_stream.hh"
+#include <stdint.h>
 
 namespace lm { namespace builder {
  
@@ -14,14 +16,18 @@ namespace lm { namespace builder {
  */
 class Interpolate {
   public:
-    explicit Interpolate(uint64_t unigram_count, const ChainPositions &backoffs);
+    // Normally vocab_size is the unigram count-1 (since p(<s>) = 0) but might
+    // be larger when the user specifies a consistent vocabulary size.
+    explicit Interpolate(uint64_t vocab_size, const util::stream::ChainPositions &backoffs, const std::vector<uint64_t> &prune_thresholds, bool output_q_);
 
-    void Run(const ChainPositions &positions);
+    void Run(const util::stream::ChainPositions &positions);
 
   private:
     float uniform_prob_;
-    ChainPositions backoffs_;
+    util::stream::ChainPositions backoffs_;
+    const std::vector<uint64_t> prune_thresholds_;
+    bool output_q_;
 };
 
 }} // namespaces
-#endif // LM_BUILDER_INTERPOLATE__
+#endif // LM_BUILDER_INTERPOLATE_H