1 files changed, 25 insertions, 3 deletions
diff --git a/lm/builder/adjust_counts.hh b/lm/builder/adjust_counts.hh
index 60198e8f8..a5435c282 100644
--- a/lm/builder/adjust_counts.hh
+++ b/lm/builder/adjust_counts.hh
@@ -2,6 +2,7 @@
 #define LM_BUILDER_ADJUST_COUNTS_H
 
 #include "lm/builder/discount.hh"
+#include "lm/lm_exception.hh"
 #include "util/exception.hh"
 
 #include <vector>
@@ -19,6 +20,16 @@ class BadDiscountException : public util::Exception {
     ~BadDiscountException() throw();
 };
 
+struct DiscountConfig {
+  // Overrides discounts for orders [1,discount_override.size()].
+  std::vector<Discount> overwrite;
+  // If discounting fails for an order, copy them from here.
+  Discount fallback;
+  // What to do when discounts are out of range or would trigger divison by
+  // zero.  It it does something other than THROW_UP, use fallback_discount.
+  WarningAction bad_action;
+};
+
 /* Compute adjusted counts.  
  * Input: unique suffix sorted N-grams (and just the N-grams) with raw counts.
  * Output: [1,N]-grams with adjusted counts.  
@@ -27,17 +38,28 @@ class BadDiscountException : public util::Exception {
  */
 class AdjustCounts {
   public:
-    AdjustCounts(std::vector<uint64_t> &counts, std::vector<uint64_t> &counts_pruned, std::vector<Discount> &discounts, std::vector<uint64_t> &prune_thresholds)
-      : counts_(counts), counts_pruned_(counts_pruned), discounts_(discounts), prune_thresholds_(prune_thresholds)
+    // counts: output
+    // counts_pruned: output
+    // discounts: mostly output.  If the input already has entries, they will be kept.
+    // prune_thresholds: input.  n-grams with normal (not adjusted) count below this will be pruned.
+    AdjustCounts(
+        const std::vector<uint64_t> &prune_thresholds,
+        std::vector<uint64_t> &counts,
+        std::vector<uint64_t> &counts_pruned,
+        const DiscountConfig &discount_config,
+        std::vector<Discount> &discounts)
+      : prune_thresholds_(prune_thresholds), counts_(counts), counts_pruned_(counts_pruned), discount_config_(discount_config), discounts_(discounts) 
     {}
 
     void Run(const util::stream::ChainPositions &positions);
 
   private:
+    const std::vector<uint64_t> &prune_thresholds_; 
     std::vector<uint64_t> &counts_;
     std::vector<uint64_t> &counts_pruned_;
+
+    DiscountConfig discount_config_;
     std::vector<Discount> &discounts_;
-    std::vector<uint64_t> &prune_thresholds_; 
 };
 
 } // namespace builder