react-native-executorch 0.10.6 → 0.10.7

package/third-party/common/phonemis/src/phonemis/base/processor/layers/term2word/layer.cpp

@@ -0,0 +1,170 @@
+#include "../../common/constants.h"
+#include "../../common/text.h"
+#include "layer.h"
+#include <phonemis/utils/strings.h>
+#include <phonemis/utils/unicode.h>
+
+#include <vector>
+
+namespace phonemis::processor::term2word {
+
+using namespace utils;
+
+namespace {
+// Whether a host label is a top-level domain, which makes a bare host an address: example.com
+bool is_domain(std::u32string_view label) {
+  static const std::unordered_set<std::u32string> domains = {U"com", U"org", U"net",  U"edu", U"gov", U"io", U"pl",
+                                                             U"de",  U"fr",  U"es",   U"it",  U"pt",  U"uk", U"eu",
+                                                             U"at",  U"ch",  U"info", U"tv",  U"co",  U"br"};
+  return domains.contains(strings::to_lower(label));
+}
+
+// A character of the user name of an email: jan.kowalski+news
+bool is_user(char32_t c) {
+  return unicode::isword(c) || strings::contains(std::u32string_view(U".-+%"), c);
+}
+
+// A character of a host label: my-site
+bool is_label(char32_t c) {
+  return unicode::isword(c) || c == U'-';
+}
+
+// A part of an address with its dots read: jan.kowalski -> jan dot kowalski
+std::u32string dotted(std::u32string_view part, const std::u32string& dot) {
+  std::u32string words;
+  for (char32_t c : part) {
+    words += c == U'.' ? U" " + dot + U" " : std::u32string(1, c);
+  }
+  return words;
+}
+} // namespace
+
+Term2WordLayer::Term2WordLayer(Terms terms) : terms_(std::move(terms)) {
+  for (const auto& [abbreviation, readings] : terms_.abbreviations) {
+    keys_.push_back(abbreviation);
+  }
+}
+
+void Term2WordLayer::apply(Document& document) const {
+  std::u32string_view text = document.text();
+
+  for (size_t position = 0; position < text.size();) {
+    // Terms start at word boundaries, but not after a dot: e.g. is not g.
+    std::optional<Term> term;
+    if (is_clear_before(text, position, U".")) {
+      term = read_address(text, position);
+      if (!term) {
+        term = read_abbreviation(text, position);
+      }
+    }
+
+    if (!term) {
+      position++;
+      continue;
+    }
+    document.replace(term->span, std::move(term->words));
+    position = term->span.end;
+  }
+}
+
+std::optional<Term2WordLayer::Term> Term2WordLayer::read_address(std::u32string_view text, size_t position) const {
+  const WebWords& web = terms_.web;
+
+  // 1. An optional scheme, which is not read: https://
+  size_t host = position;
+  while (host < text.size() && unicode::is_latin(text[host])) {
+    host++;
+  }
+  bool scheme = host > position && text.substr(host).starts_with(U"://");
+  host = scheme ? host + 3 : position;
+
+  // 2. The user of an email: jan.kowalski@
+  std::u32string_view user;
+  if (size_t at = host; !scheme) {
+    while (at < text.size() && is_user(text[at])) {
+      at++;
+    }
+    if (at > host && at < text.size() && text[at] == U'@') {
+      user = text.substr(host, at - host);
+      host = at + 1;
+    }
+  }
+
+  // 3. The host: labels joined by dots
+  std::vector<std::u32string_view> labels;
+  size_t end = host;
+  while (true) {
+    size_t label_end = end;
+    while (label_end < text.size() && is_label(text[label_end])) {
+      label_end++;
+    }
+    if (label_end == end) {
+      break;
+    }
+    labels.push_back(text.substr(end, label_end - end));
+    end = label_end;
+    if (end + 1 >= text.size() || text[end] != U'.' || !is_label(text[end + 1])) {
+      break;
+    }
+    end++;
+  }
+
+  // A bare host needs a known top-level domain: example.com, but not file.txt or e.g
+  bool www = !labels.empty() && strings::to_lower(labels[0]) == U"www";
+  bool domain = labels.size() >= 2 && is_domain(labels.back());
+  if (labels.size() < 2 || !(scheme || www || domain || !user.empty())) {
+    return std::nullopt;
+  }
+
+  // 4. An optional path, without the punctuation closing it: (see example.com/news).
+  size_t path_end = end;
+  if (end < text.size() && text[end] == U'/') {
+    while (path_end < text.size() && !unicode::isspace(text[path_end])) {
+      path_end++;
+    }
+    while (path_end > end + 1 && strings::contains(constants::kPunctuation, text[path_end - 1])) {
+      path_end--;
+    }
+  }
+
+  // 5. The words: jan at example dot com, w w w dot example dot com slash news
+  std::u32string words = user.empty() ? U"" : dotted(user, web.dot) + U" " + web.at + U" ";
+  for (size_t i = 0; i < labels.size(); ++i) {
+    bool spelled = i == 0 && www;
+    words += (i > 0 ? U" " + web.dot + U" " : U"") + (spelled ? U"w w w" : std::u32string(labels[i]));
+  }
+  for (auto path = text.substr(end, path_end - end); !path.empty();) {
+    size_t next = path.find(U'/', 1);
+    std::u32string_view segment = path.substr(1, next == std::u32string_view::npos ? next : next - 1);
+    words += U" " + web.slash + (segment.empty() ? U"" : U" " + dotted(segment, web.dot));
+    path = next == std::u32string_view::npos ? U"" : path.substr(next);
+  }
+  return Term{.span = {position, path_end}, .words = words};
+}
+
+std::optional<Term2WordLayer::Term> Term2WordLayer::read_abbreviation(std::u32string_view text, size_t position) const {
+  for (auto key : keys_at(text, position, keys_)) {
+    // An abbreviation ending with a letter ends a word, or touches a number: nr5
+    size_t end = position + key.size();
+    if (unicode::isword(key.back()) && !is_clear_after(text, end) && !unicode::isdigit(text[end])) {
+      continue;
+    }
+
+    auto words = read(terms_.abbreviations.at(std::u32string(key)), Context(text, {position, end}));
+    if (!words) {
+      continue;
+    }
+
+    // A dot which also ends the text stays with it; a word right after gets a space: nr5 -> numer 5
+    if (key.ends_with(U'.') && is_sentence_end(text, end)) {
+      *words += U".";
+    }
+    if (end < text.size() && unicode::isword(text[end])) {
+      *words += U" ";
+    }
+    return Term{.span = {position, end}, .words = *words};
+  }
+  return std::nullopt;
+}
+
+} // namespace phonemis::processor::term2word

package/third-party/common/phonemis/src/phonemis/base/processor/layers/term2word/layer.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include "../../layer.h"
+#include "types.h"
+
+#include <optional>
+#include <string>
+#include <string_view>
+
+namespace phonemis::processor::term2word {
+
+/**
+ * Handles non-numeric and non-symbolic terms.
+ * 
+ * Example: web addresses (www.example.com/news, jan@example.com), abbreviations (nr 5, e.g.).
+ */
+class Term2WordLayer : public Layer {
+public:
+  explicit Term2WordLayer(Terms terms);
+
+  void apply(Document& document) const override;
+
+  // A term found in a text, with its words
+  struct Term {
+    Span span;
+    std::u32string words;
+  };
+
+private:
+  // https://www.example.com/news, jan.kowalski@example.pl: its parts joined with the words of the separators
+  std::optional<Term> read_address(std::u32string_view text, size_t position) const;
+
+  // The longest abbreviation at a position with a reading in its context
+  std::optional<Term> read_abbreviation(std::u32string_view text, size_t position) const;
+
+  Terms terms_;
+  Keys keys_; // Of the abbreviations
+};
+
+} // namespace phonemis::processor::term2word

package/third-party/common/phonemis/src/phonemis/base/processor/layers/term2word/types.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include "../../common/conditions.h"
+
+#include <string>
+
+namespace phonemis::processor::term2word {
+
+// Abbreviations with their readings: nr -> numer before a number; St. -> Saint before a name, Street otherwise.
+using Abbreviations = Readings;
+
+// The words between the parts of a web address.
+struct WebWords {
+  std::u32string dot;   // example dot com
+  std::u32string at;    // jan at example dot com
+  std::u32string slash; // example dot com slash news
+};
+
+// The terms of a language: written forms which are not said as written.
+struct Terms {
+  Abbreviations abbreviations;
+  WebWords web;
+};
+
+} // namespace phonemis::processor::term2word

package/third-party/common/phonemis/src/phonemis/base/processor/layers/trim/layer.cpp

@@ -0,0 +1,47 @@
+#include "layer.h"
+#include <phonemis/utils/strings.h>
+#include <phonemis/utils/unicode.h>
+
+namespace phonemis::processor::trim {
+
+namespace {
+// Whether a position starts punctuation which takes no space before it: "a ," -> "a,", but "a ..." stays
+bool is_closing_punctuation(std::u32string_view text, size_t position) {
+  if (position >= text.size()) {
+    return false;
+  }
+  bool ellipsis = text.substr(position).starts_with(U"...");
+  return utils::strings::contains(std::u32string_view(U",!?;:"), text[position]) ||
+         (text[position] == U'.' && !ellipsis);
+}
+} // namespace
+
+void TrimLayer::apply(Document& document) const {
+  std::u32string_view text = document.text();
+
+  for (size_t i = 0; i < text.size();) {
+    if (!utils::unicode::isspace(text[i])) {
+      i++;
+      continue;
+    }
+
+    // Find the whole run of whitespaces
+    size_t end = i;
+    bool newline = false;
+    while (end < text.size() && utils::unicode::isspace(text[end])) {
+      newline |= text[end] == U'\n';
+      end++;
+    }
+
+    // Edges and spaces before punctuation lose the run entirely, the middle keeps a single character
+    bool closing = !newline && is_closing_punctuation(text, end);
+    std::u32string replacement = i == 0 || end == text.size() || closing ? U"" : newline ? U"\n" : U" ";
+    if (text.substr(i, end - i) != replacement) {
+      document.replace({i, end}, replacement);
+    }
+
+    i = end;
+  }
+}
+
+} // namespace phonemis::processor::trim

package/third-party/common/phonemis/src/phonemis/base/processor/layers/trim/layer.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../../layer.h"
+
+namespace phonemis::processor::trim {
+
+/**
+ * Normalizes whitespaces: removes them from both ends of the text and before punctuation (a , -> a,),
+ * and collapses every other run into a single "\n" (if it contains a line break) or " ".
+ */
+class TrimLayer : public Layer {
+public:
+  TrimLayer() = default;
+  ~TrimLayer() override = default;
+
+  void apply(Document& document) const override;
+};
+
+} // namespace phonemis::processor::trim

package/third-party/common/phonemis/src/phonemis/base/processor/preprocessor.cpp

@@ -0,0 +1,60 @@
+#include "layers/normalize/layer.h"
+#include "layers/term2word/layer.h"
+#include "layers/trim/layer.h"
+#include "preprocessor.h"
+
+namespace phonemis::processor {
+
+namespace {
+// A table with extra entries, which replace those with the same key
+Readings merged(Readings table, const Readings& extra) {
+  for (const auto& [key, readings] : extra) {
+    table.insert_or_assign(key, readings);
+  }
+  return table;
+}
+} // namespace
+
+Preprocessor make_preprocessor(const Config& config, const Language& language, std::unique_ptr<Layer> numbers) {
+  Preprocessor preprocessor;
+
+  // Characters normalized first, so that every layer sees one form of each
+  if (config.normalize) {
+    preprocessor.add_layer(std::make_unique<normalize::NormalizeLayer>(
+        language.characters ? *language.characters : normalize::constants::kCharacters));
+  }
+
+  if (config.trim) {
+    preprocessor.add_layer(std::make_unique<trim::TrimLayer>());
+  }
+
+  // Terms read before numbers, so that "nr 5" still sees its digit
+  if (config.terms && (language.terms || !config.abbreviations.empty())) {
+    term2word::Terms terms = language.terms ? *language.terms : term2word::Terms{};
+    terms.abbreviations = merged(std::move(terms.abbreviations), config.abbreviations);
+    preprocessor.add_layer(std::make_unique<term2word::Term2WordLayer>(std::move(terms)));
+  }
+
+  if (numbers) {
+    preprocessor.add_layer(std::move(numbers));
+  }
+
+  // Symbols read after numbers, which read the ones they own: 20%, $5
+  if (config.symbols && (language.symbols || !config.symbol_words.empty())) {
+    sym2word::Symbols symbols = language.symbols ? *language.symbols : sym2word::Symbols{};
+    preprocessor.add_layer(std::make_unique<sym2word::Sym2WordLayer>(merged(std::move(symbols), config.symbol_words)));
+  }
+
+  if (config.sanitize && language.replacements) {
+    const Replacements& replacements = *language.replacements;
+    auto mapper = [&replacements](char32_t c) {
+      auto it = replacements.find(c);
+      return it != replacements.end() ? it->second : c;
+    };
+    preprocessor.add_layer(std::make_unique<sanitizer::SanitizerLayer>([](char32_t) { return true; }, mapper));
+  }
+
+  return preprocessor;
+}
+
+} // namespace phonemis::processor

package/third-party/common/phonemis/src/phonemis/base/processor/preprocessor.h

@@ -0,0 +1,45 @@
+#pragma once
+
+#include "config.h"
+#include "layers/normalize/types.h"
+#include "layers/sanitizer/layer.h"
+#include "layers/sym2word/layer.h"
+#include "layers/term2word/types.h"
+#include "processor.h"
+
+#include <concepts>
+#include <memory>
+#include <unordered_map>
+
+namespace phonemis::processor {
+
+// Characters replaced by others before phonemization.
+using Replacements = std::unordered_map<char32_t, char32_t>;
+
+// The parts of a language's preprocessing, from its constants.
+struct Language {
+  const normalize::Characters* characters = nullptr; // Its normalized characters, if not the default ones
+  const term2word::Terms* terms = nullptr;           // Its abbreviations and web words, if any
+  const sym2word::Symbols* symbols = nullptr;        // Its symbol words, if any
+  const Replacements* replacements = nullptr;        // Its character replacements, if any: ’ -> '
+};
+
+/**
+ * Builds a language's preprocessing from the layers its configuration enables, in their fixed order:
+ * Normalize, Trim, Term2Word, Num2Word, Sym2Word, Sanitizer.
+ * @param config which layers run, and words added to the language's tables.
+ * @param language the language's tables.
+ * @param numbers the language's Num2Word layer, or nullptr if numbers are not read.
+ */
+Preprocessor make_preprocessor(const Config& config, const Language& language, std::unique_ptr<Layer> numbers);
+
+/**
+ * Builds a language's preprocessing, as above.
+ * @tparam N2W the language's Num2Word layer, created only if numbers are read.
+ */
+template <std::derived_from<Layer> N2W>
+Preprocessor make_preprocessor(const Config& config, const Language& language) {
+  return make_preprocessor(config, language, config.numbers ? std::make_unique<N2W>() : nullptr);
+}
+
+} // namespace phonemis::processor

package/third-party/common/phonemis/src/phonemis/base/processor/processor.cpp

@@ -1,3 +1,4 @@
+#include "common/document.h"
 #include "processor.h"
 
 namespace phonemis::processor {
@@ -7,17 +8,15 @@ void Processor::add_layer(std::unique_ptr<Layer> layer) {
 }
 
 std::u32string Processor::process(std::u32string_view input) const {
-  if (layers_.empty()) {
-    return std::u32string(input);
-  }
-
-  std::u32string result(input);
+  // We start by initializing a document - an object to work for all the layers.
+  Document document(input);
 
   for (const auto& layer : layers_) {
-    result = layer->transform(result);
+    layer->apply(document);
+    document.commit();
   }
 
-  return result;
+  return std::u32string(document.text());
 }
 
 } // namespace phonemis::processor

package/third-party/common/phonemis/src/phonemis/base/processor/processor.h

@@ -7,27 +7,31 @@
 
 namespace phonemis::processor {
 
-// Can be customized in derived classes by adding different preprocessing layers.
+/**
+ * Main preprocessing component.
+ * 
+ * A composition of elementary steps - can be customized by specyfing different layers (steps).
+ */
 class Processor {
 public:
   Processor() = default;
 
   /**
-   * Allows to customize Processor's behavior by expanding it with custom layers.
-   * @param layer a new transformation layer to be added.
-   */
-  void add_layer(std::unique_ptr<Layer> layer);
-
-  /**
-   * Processes the input by applying all of it's transformation layers.
+   * Processes the input by applying all of it's transformation layers, one after another.
+   * @details Everything the layers learn about the text lives in a Document, which is dropped on return.
    * @param input an input to be processed
    * @returns processed input.
    */
   std::u32string process(std::u32string_view input) const;
 
+  /**
+   * Allows to customize Processor's behavior by expanding it with custom layers.
+   * @param layer a new transformation layer to be added.
+   */
+  void add_layer(std::unique_ptr<Layer> layer);
+
 private:
-  // Preprocessing layers - each layer performs some sort of
-  // text -> text transformation on its input.
+  // Preprocessing layers - each layer queues edits of a shared document.
   std::vector<std::unique_ptr<Layer>> layers_;
 };
 

package/third-party/common/phonemis/src/phonemis/base/tagger/hmm_tagger.cpp

@@ -13,6 +13,11 @@ namespace phonemis::tagger {
 
 using namespace utils;
 
+namespace {
+// Small constant for unseen observations/transitions in log space.
+const double LOG_EPSILON = std::log(1e-7);
+}  // namespace
+
 HMMTagger::HMMTagger(const Config& config) {
   if (!config.data_filepath.has_value()) {
     throw std::runtime_error("HMMTagger: data_filepath must be provided in the configuration.");
@@ -28,95 +33,127 @@ HMMTagger::HMMTagger(const Config& config) {
     }
   }
 
-  // 1. Start Probabilities (defines the tag set)
-  // Store as log probabilities (numeric stability)
+  // Parses a JSON probability into log-space, rejecting values outside [0, 1]
+  // (including negatives and NaN, which would make std::log produce NaN and
+  // silently corrupt the Viterbi computation). A stored 0 maps to -inf, which
+  // is valid in log space (an impossible event).
+  auto to_log_prob = [](const nlohmann::json& node, const std::string& context) -> double {
+    const double p = node.get<double>();
+    if (!(p >= 0.0 && p <= 1.0)) {
+      throw std::invalid_argument(
+          "HMMTagger: probability for '" + context + "' must be in [0, 1], got " + std::to_string(p));
+    }
+    return std::log(p);
+  };
+
+  // 1. The start-probability keys define the tag set. Assign each tag a dense
+  //    integer id; sorting keeps the id assignment (and thus Viterbi
+  //    tie-breaking) deterministic across runs.
+  for (const auto& [tag, prob] : json["start_prob"].items()) {
+    tag_names_.push_back(tag);
+  }
+  std::ranges::sort(tag_names_);
+
+  const size_t T = tag_names_.size();
+  std::unordered_map<Tag, size_t> tag_id;
+  tag_id.reserve(T);
+  for (size_t id = 0; id < T; ++id) {
+    tag_id.emplace(tag_names_[id], id);
+  }
+
+  // 2. Start probabilities, indexed by tag id (stored as log probabilities).
+  start_probs_.assign(T, LOG_EPSILON);
   for (const auto& [tag, prob] : json["start_prob"].items()) {
-    tags_.insert(tag);
-    start_probs_[tag] = std::log(prob.get<double>());
+    start_probs_[tag_id.at(tag)] = to_log_prob(prob, tag);
   }
 
-  // 2. Emission Probabilities: P(Word | Tag)
-  // Store as log probabilities (numeric stability)
+  // 3. Emission probabilities: P(Word | Tag), one word map per tag id. Tags
+  //    not present in the start set (and therefore not in the model's state
+  //    space) are ignored, mirroring the original behavior.
+  emission_probs_.resize(T);
   for (const auto& [tag, words] : json["emission"].items()) {
+    auto it = tag_id.find(tag);
+    if (it == tag_id.end()) continue;
     for (const auto& [word, prob] : words.items()) {
-      emission_probs_[tag][word] = std::log(prob.get<double>());
+      emission_probs_[it->second][word] = to_log_prob(prob, tag + "/" + word);
     }
   }
 
-  // 3. Transition Probabilities: P(Tag_i | Tag_{i-1})
-  // Store as log probabilities (numeric stability)
+  // 4. Transition probabilities: P(Tag_i | Tag_{i-1}) as a flat T x T matrix,
+  //    pre-filled with LOG_EPSILON so missing transitions need no runtime check.
+  transition_probs_.assign(T * T, LOG_EPSILON);
   for (const auto& [prev_tag, next_tags] : json["transition"].items()) {
+    auto prev_it = tag_id.find(prev_tag);
+    if (prev_it == tag_id.end()) continue;
     for (const auto& [next_tag, prob] : next_tags.items()) {
-      transition_probs_[prev_tag][next_tag] = std::log(prob.get<double>());
+      auto next_it = tag_id.find(next_tag);
+      if (next_it == tag_id.end()) continue;
+      transition_probs_[prev_it->second * T + next_it->second] =
+          to_log_prob(prob, prev_tag + "->" + next_tag);
     }
   }
-
-  // Ensure all tags exist in emission and transition maps
-  for (const auto& tag : tags_) {
-    emission_probs_[tag];    // Ensures outer map entry exists
-    transition_probs_[tag];  // Ensures outer map entry exists
-  }
 }
 
 void HMMTagger::tag_sentence(std::span<Token> sentence) const {
   if (sentence.empty()) return;
 
-  // Small constant for unseen observations in log space
-  const double LOG_EPSILON = std::log(1e-7);
+  const size_t T = tag_names_.size();
+  if (T == 0) return;  // No model loaded -> nothing to tag.
+
+  const size_t n = sentence.size();
+
+  // The recursion only ever reads the previous column, so dp keeps just two
+  // rolling columns. backtrace keeps the full [step][state] grid of tag ids.
+  std::vector<double> dp_prev(T), dp_curr(T);
+  std::vector<std::vector<int>> backtrace(n, std::vector<int>(T, 0));
 
-  // Viterbi tables: [step][state]
-  std::vector<std::unordered_map<Tag, double>> dp(sentence.size());
-  std::vector<std::unordered_map<Tag, Tag>> backtrace(sentence.size());
+  // log P(word | tag id), falling back to LOG_EPSILON for unseen words.
+  auto emit_log_prob = [this](size_t tag_id, const std::string& word) {
+    const auto& word_map = emission_probs_[tag_id];
+    auto it = word_map.find(word);
+    return it != word_map.end() ? it->second : LOG_EPSILON;
+  };
 
   // Step 0: Initialization
   std::string first_word = conversions::u32_to_utf8(
-    should_lowercase(sentence[0].text)    ? 
-      strings::to_lower(sentence[0].text) : 
+    should_lowercase(sentence[0].text) ?
+      strings::to_lower(sentence[0].text) :
       sentence[0].text
   );
-
-  for (const auto& tag : tags_) {
-    double emit_p = emission_probs_.at(tag).contains(first_word) 
-                    ? emission_probs_.at(tag).at(first_word) : LOG_EPSILON;
-    dp[0][tag] = start_probs_.at(tag) + emit_p;
+  for (size_t tag = 0; tag < T; ++tag) {
+    dp_prev[tag] = start_probs_[tag] + emit_log_prob(tag, first_word);
   }
 
-  // Steps 1..N: Viterbi Recursion
-  for (size_t t = 1; t < sentence.size(); ++t) {
+  // Steps 1..n-1: Viterbi Recursion
+  for (size_t t = 1; t < n; ++t) {
     std::string word = conversions::u32_to_utf8(sentence[t].text);
 
-    for (const auto& curr_tag : tags_) {
-      double max_p = -std::numeric_limits<double>::infinity();
-      Tag best_prev = *tags_.begin();
-
-      double emit_p = emission_probs_.at(curr_tag).contains(word) 
-                      ? emission_probs_.at(curr_tag).at(word) : LOG_EPSILON;
+    for (size_t curr = 0; curr < T; ++curr) {
+      const double emit_p = emit_log_prob(curr, word);
 
-      for (const auto& prev_tag : tags_) {
-        double trans_p = transition_probs_.at(prev_tag).contains(curr_tag)
-                         ? transition_probs_.at(prev_tag).at(curr_tag) : LOG_EPSILON;
-        
-        double p = dp[t - 1][prev_tag] + trans_p + emit_p;
+      double max_p = -std::numeric_limits<double>::infinity();
+      int best_prev = 0;
+      for (size_t prev = 0; prev < T; ++prev) {
+        const double p = dp_prev[prev] + transition_probs_[prev * T + curr] + emit_p;
         if (p > max_p) {
           max_p = p;
-          best_prev = prev_tag;
+          best_prev = static_cast<int>(prev);
         }
       }
-      dp[t][curr_tag] = max_p;
-      backtrace[t][curr_tag] = best_prev;
+      dp_curr[curr] = max_p;
+      backtrace[t][curr] = best_prev;
     }
+    std::swap(dp_prev, dp_curr);  // dp_prev now holds the column for step t.
   }
 
-  // Backtracking: Find most likely end state and reverse-walk the path
-  auto last_idx = sentence.size() - 1;
-  Tag current_tag = *std::max_element(tags_.begin(), tags_.end(), [&](const Tag& a, const Tag& b) {
-    return dp[last_idx][a] < dp[last_idx][b];
-  });
+  // Backtracking: pick the most likely final state (dp_prev holds the last
+  // column after the final swap), then reverse-walk the path.
+  int current = static_cast<int>(std::ranges::max_element(dp_prev) - dp_prev.begin());
 
-  for (size_t t = last_idx; ; --t) {
-    sentence[t].tag = current_tag;
+  for (size_t t = n - 1; ; --t) {
+    sentence[t].tag = tag_names_[current];
     if (t == 0) break;
-    current_tag = backtrace[t][current_tag];
+    current = backtrace[t][current];
   }
 }