dynamic pruning and substitution restrictions (#60)

* massive speed improvement, with some WER improvements too in large complicated files with lots of deletions
* adding a flexible beam, adjustable via command line parameters
* adding --strict-punctuation mode, that will allow punctuation marks to be substituted only within themselves
* adding support for strict punctuation and favouring same words in alignments
* fixing the default value of the beam to 50
* adding new test case file
* adding new result for std composition
* fixes, and new test cases updates
* bumping version and adding release notes to the readme

Author: jprobichaud

README.md

@@ -3,6 +3,7 @@
 
 # fstalign
 - [Overview](#Overview)
+- [What's new in 2.0](#What's-new-in-2.0)
 - [Installation](#Installation)
   * [Dependencies](#Dependencies)
   * [Build](#Build)
@@ -14,6 +15,97 @@
 
 Due to its use of OpenFST and lazy algorithms for text-based alignment, `fstalign` is efficient for calculating WER while also providing significant flexibility for different measurement features and error analysis.
 
+## What's new in 2.0
+
+Version 2.0 introduces two major changes: 
+1. A new method to traverse the composition graph, which dramatically improves the overall speed, especially when the sequences are long contain many errors.
+We have files that took 25 minutes to align before that can now take about 7 seconds. This is especially noticeable with the adapted composition (the default).
+1. Some smarts were introduced when --use-case and --use-punctuation are enabled.
+Now, by default, punctuation symbols can only be substituted by other punctuation symbols (or deleted/inserted).
+Also, words that differ only by the first letter case will be preffered for substitution.
+
+
+Here's an example of the 1.x behavior and the 2.0 version
+```
+==> v1.x sbs.txt <==
+           ref_token	hyp_token           	IsErr	Class	Wer_Tag_Entities
+             Welcome	Welcome             			###322_###|
+                back	back                			
+                  to	to                  			
+             another	another             			
+             episode	episode             			###323_###|
+                  of	of                  			
+            Podcasts	Podcast             	ERR		###324_###|
+                  in	and                 	ERR		
+               Color	Color               			###167_###|###325_###|
+                   :	of                  	ERR		
+                 The	the                 	ERR		
+             Podcast	Podcast             			###168_###|###326_###|
+                   .	.                   			
+                   I	I                   			
+
+==> v2.0 sbs.txt <==
+           ref_token	hyp_token           	IsErr	Class	Wer_Tag_Entities
+             Welcome	Welcome             			###322_###|
+                back	back                			
+                  to	to                  			
+             another	another             			
+             episode	episode             			###323_###|
+                  of	of                  			
+            Podcasts	Podcast             	ERR		###324_###|
+                  in	and                 	ERR		
+               Color	Color               			###167_###|###325_###|
+               <ins>	of                  	ERR		
+                   :	<del>               	ERR		
+                 The	the                 	ERR		
+             Podcast	Podcast             			###168_###|###326_###|
+```
+The confusion between `:` and `of` is not longer allowed.
+
+Also, here's how favoring or not the substitution based on case-insensitive comparison, while still counting it as an error, looks like:
+```
+==> v1.x sbs.txt <==
+           ref_token	hyp_token           	IsErr	Class	Wer_Tag_Entities
+             shorten    shorten                         ###801_###|
+                It's    it's                    ERR     
+               Berry    Barry                   ERR     ###785_###|###788_###|###802_###|
+                   .    .        
+                Just    Just   
+                Yeah    like                    ERR     ###805_###|                                                                                                             
+                   .    <del>                   ERR
+                Like    <del>                   ERR     
+                   ,    <del>                   ERR     
+                   I    I                               ###809_###|
+                have    have   
+                   a    a         
+            nickname    nickname 
+
+==> v2.0 sbs.txt <==
+           ref_token	hyp_token           	IsErr	Class	Wer_Tag_Entities
+                It's    it's                    ERR     
+               Berry    Barry                   ERR     ###785_###|###788_###|###802_###|
+                   .    .     
+                Just    Just     
+                Yeah    <del>                   ERR     ###805_###|
+                   .    <del>                   ERR     
+                Like    like                    ERR     
+                   ,    <del>                   ERR     
+                   I    I                               ###809_###|
+                have    have     
+                   a    a        
+            nickname    nickname  
+```
+Here, `Like <-> like` substitution is favored.  While this generally won't change the WER value itself (although it can), it will improve the timing alignments.  
+
+
+These behavior, as well as the beam size (that has a default value of 50.0) can be controlled with the following new parameters:
+```
+  --disable-strict-punctuation
+                              Disable strict punctuation alignment (which prevents punctuation aligning with words).
+  --disable-favored-subs      Disable favored substitutions (which makes alignment favor substitutions between words which differ only by case).
+  --favored-sub-cost FLOAT    Cost for favored substitutions (e.g., case diff). Default: 0.1
+```
+
 ## Installation
 
 ### Dependencies

src/AdaptedComposition.cpp

@@ -21,7 +21,7 @@ AdaptedCompositionFst::AdaptedCompositionFst(const fst::StdFst &fstA, const fst:
 #endif
 }
 
-AdaptedCompositionFst::AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, SymbolTable &symbols)
+AdaptedCompositionFst::AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, const SymbolTable &symbols)
     : fstA_{fstA}, fstB_{fstB} {
   logger_ = logger::GetOrCreateLogger("AdaptedCompositionFst");
   logger_->set_level(spdlog::level::info);
@@ -36,6 +36,35 @@ AdaptedCompositionFst::AdaptedCompositionFst(const fst::StdFst &fstA, const fst:
   ins_label_id_ = symbols.Find(options.symIns);
 }
 
+AdaptedCompositionFst::AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, const SymbolTable &symbols, const AlignerOptions& options)
+    : fstA_(fstA),                     // Reference to input FST A
+      fstB_(fstB),                     // Reference to input FST B
+      symbols_(&symbols),              // Pointer to symbol table
+      strict_punctuation_(options.strict_punctuation), // Store strict punctuation flag
+      punctuation_ids_(options.punctuation_ids),
+      // Favored substitutions
+      use_favored_substitutions_(options.use_favored_substitutions),
+      favored_substitution_cost_(options.favored_substitution_cost),
+      favorable_substitution_map_(options.favorable_substitution_map)
+      // Initialize other members if they exist (e.g., current_composed_next_state_id = 0;)
+{
+    logger_ = logger::GetOrCreateLogger("AdaptedCompositionFst"); // Use member logger_ if declared
+    logger_->set_level(spdlog::level::info);
+    #if TRACE
+      logger_->set_level(spdlog::level::trace);
+    #endif
+
+    // Initialize special symbol IDs (assuming these are member variables now)
+    FstAlignOption fst_options; // Contains special symbol names (symSub etc)
+    // Ensure sub_label_id_ etc are members if used elsewhere
+    sub_label_id_ = symbols_->Find(fst_options.symSub);
+    del_label_id_ = symbols_->Find(fst_options.symDel);
+    ins_label_id_ = symbols_->Find(fst_options.symIns);
+
+    // Initialize entity/synonym label vectors
+    SetSymbols(&symbols);
+}
+
 AdaptedCompositionFst::~AdaptedCompositionFst() {}
 
 bool AdaptedCompositionFst::IsEntityLabel(int labelId) {
@@ -349,8 +378,40 @@ bool AdaptedCompositionFst::TryGetArcsAtState(StateId fromStateId, vector<fst::S
 #if TRACE
         logger_->trace("{}/{} >] adding sub/{}/{}", dbg_count, here_snap, arcA.ilabel, arcB.olabel);
 #endif
-        out_vector->push_back(StdArc(arcA.ilabel, arcB.olabel, substitution_cost, sub_state_ref_id));
-        arc_added++;
+        // --- Strict Punctuation Check --- 
+        bool skip_substitution = false;
+        if (strict_punctuation_) {
+            bool ilabel_is_punct = (punctuation_ids_.count(arcA.ilabel) > 0);
+            bool olabel_is_punct = (punctuation_ids_.count(arcB.olabel) > 0); // Check arcB.olabel for hyp side?
+             
+            if (ilabel_is_punct != olabel_is_punct) {
+                skip_substitution = true;
+            }
+        }
+        // --- End Strict Punctuation Check ---
+        
+        if (!skip_substitution) {
+          // --- Favored Substitution Cost Check ---
+          float current_sub_cost = 1.0f; // Default substitution cost
+          if (use_favored_substitutions_) {
+              int labelA = arcA.ilabel;
+              int labelB = arcB.olabel;
+              // Check bounds and if labelA has a favored partner which is labelB
+              if (labelA >= 0 && labelA < favorable_substitution_map_.size() &&
+                  favorable_substitution_map_[labelA] == labelB)
+              {
+                  current_sub_cost = favored_substitution_cost_; // Use lower cost
+#if TRACE
+                   logger_->trace("Applying favored sub cost ({}) for {} ({}) <-> {} ({})",
+                                 current_sub_cost, symbols_->Find(labelA), labelA, symbols_->Find(labelB), labelB);
+#endif
+              }
+          }
+          // --- End Favored Substitution Cost Check ---
+
+          out_vector->push_back(StdArc(arcA.ilabel, arcB.olabel, current_sub_cost, sub_state_ref_id));
+          arc_added++;
+        }
       }
     }
 
@@ -402,7 +463,7 @@ bool AdaptedCompositionFst::TryGetArcsAtState(StateId fromStateId, vector<fst::S
   return true;
 }
 
-void AdaptedCompositionFst::SetSymbols(fst::SymbolTable *symbols) {
+void AdaptedCompositionFst::SetSymbols(const fst::SymbolTable *symbols) {
   symbols_ = symbols;
   synonyms_label_ids.clear();
   synonyms_label_ids.resize(symbols->NumSymbols(), false);

src/AdaptedComposition.h

@@ -13,6 +13,7 @@ AdaptedComposition.h
 #include <utility>
 #include "IComposition.h"
 #include "utilities.h"
+#include "fstalign.h"
 
 using namespace std;
 
@@ -48,7 +49,7 @@ class AdaptedCompositionFst : public IComposition {
 
   StateId current_composed_next_state_id = 0;
 
-  fst::SymbolTable *symbols_;
+  const fst::SymbolTable *symbols_;
   std::vector<bool> synonyms_label_ids;
   std::vector<bool> entity_label_ids;
 
@@ -57,6 +58,13 @@ class AdaptedCompositionFst : public IComposition {
   // possible optimizations : limit to const FST or limit to StdVectorFst
   const fst::StdFst &fstA_;
   const fst::StdFst &fstB_;
+  // Add members to store options
+  bool strict_punctuation_ = false;
+  std::unordered_set<int> punctuation_ids_;
+  // Favored substitutions
+  bool use_favored_substitutions_ = false;
+  float favored_substitution_cost_ = 0.1f;
+  std::vector<int> favorable_substitution_map_;
 
   StateId GetOrCreateComposedState(StateId a, StateId b);
   bool IsEntityLabel(int labelId);
@@ -65,7 +73,8 @@ class AdaptedCompositionFst : public IComposition {
 
  public:
   AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB);
-  AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, SymbolTable &symbols);
+  AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, const SymbolTable &symbols);
+  AdaptedCompositionFst(const fst::StdFst &fstA, const fst::StdFst &fstB, const SymbolTable &symbols, const AlignerOptions& options);
   ~AdaptedCompositionFst();
 
   StateId Start();
@@ -78,7 +87,7 @@ class AdaptedCompositionFst : public IComposition {
   // a is in the composed-graph referencial
   bool DoesComposedStateExist(StateId a);
 
-  void SetSymbols(fst::SymbolTable *symbols);
+  void SetSymbols(const fst::SymbolTable *symbols);
 
   void DebugComposedGraph();
 };

src/PathHeap.cpp

@@ -114,3 +114,48 @@ bool  pruneMe = (*last_wer_index)->numErrors + 20 < (*iter)->numErrors; --> seem
 
   return pruned;
 }
+
+int PathHeap::prune_relative(float beam_width) {
+    if (heap.empty()) {
+        return 0;
+    }
+
+    auto logger = logger::GetOrCreateLogger("pathheap");
+    size_t initial_size = heap.size();
+
+    // Find the best costSoFar in the current heap
+    // Note: The heap is ordered by the complex shortlistComparatorSharedPtr,
+    //       so the first element isn't necessarily the one with the lowest costSoFar.
+    //       We need to iterate to find the minimum costSoFar.
+    float best_cost = std::numeric_limits<float>::max();
+    for (const auto& entry : heap) {
+        if (entry->costSoFar < best_cost) {
+            best_cost = entry->costSoFar;
+        }
+    }
+
+    float cost_threshold = best_cost + beam_width;
+
+    logger->debug("==== Relative pruning starting (Beam: {}) =====", beam_width);
+    logger->debug("Initial size: {}, Best cost: {:.4f}, Threshold: {:.4f}",
+                  initial_size, best_cost, cost_threshold);
+
+    int pruned_count = 0;
+    auto iter = heap.begin();
+    while (iter != heap.end()) {
+        // Check if the current entry's cost exceeds the threshold
+        if ((*iter)->costSoFar > cost_threshold) {
+            // Remove the element and advance the iterator
+            iter = heap.erase(iter);
+            pruned_count++;
+        } else {
+            // Otherwise, just advance the iterator
+            ++iter;
+        }
+    }
+
+    logger->debug("After relative pruning: {} items remain ({} pruned)", heap.size(), pruned_count);
+    logger->debug("-----\n");
+
+    return pruned_count;
+}

src/PathHeap.h

@@ -64,6 +64,7 @@ class PathHeap {
   void insert(std::shared_ptr<ShortlistEntry> entry);
   shared_ptr<ShortlistEntry> removeFirst();
   int prune(int targetSz);
+  int prune_relative(float beam_width);
   int size();
   std::shared_ptr<ShortlistEntry> GetBestWerCandidate();
   int pruningErrorOffset = 20;