Another fix

Author: apsonawane

operators/tokenizer/case_encoder.h

@@ -33,6 +33,18 @@ class CaseEncoder {
   virtual ~CaseEncoder() {}
   void SetNormalizer(Normalizer normalizer) { normalizer_ = normalizer; }
 
+  // Reset all state for a new tokenization call
+  void Reset() {
+    buffer_.clear();
+    buffer_queue_.clear();
+    signature_.clear();
+    offset_ = 0;
+    dump_buffer_from_ = -1;
+    state_ = 0;
+    spans_ = 0;
+    seen_three_spans_ = false;
+  }
+
  public:
   CaseEncoder(bool remove_extra_white_space) : remove_extra_white_space_(remove_extra_white_space) {}
 
@@ -98,10 +110,6 @@ class CaseEncoder {
       buffer_.clear();
       buffer_queue_.clear();
       offset_ = 0;
-      // Reset all state for new tokenization
-      signature_.clear();
-      spans_ = 0;
-      seen_three_spans_ = false;
     }
 
     if (isUpper(sp)) {
@@ -174,9 +182,10 @@ class CaseEncoder {
   void PostProcess(std::string* normalized, std::vector<size_t>* norm_to_orig) {
     if (!seen_three_spans_) return;
 
-    // Safety check: ensure norm_to_orig has enough elements
-    if (norm_to_orig->size() < normalized->size()) {
-      return;  // Cannot safely process - sizes don't match
+    // Ensure norm_to_orig has at least as many elements as normalized
+    // Pad with zeros if needed to prevent out-of-bounds access
+    while (norm_to_orig->size() < normalized->size()) {
+      norm_to_orig->push_back(0);
     }
 
     std::string normalized_temp;

operators/tokenizer/ugm_kernels.hpp

@@ -535,9 +535,12 @@ struct SpmUgmTokenizer {
   }
 
   std::string NmtNormalize(const std::string& input) const {
+    // Reset the case encoder state before starting new normalization
+    case_encoder_->Reset();
+
     std::string normalized;
     normalized.reserve(input.size() * 3);
-    // Use a vector that tracks original positions - reserve capacity but start empty
+    // Track norm_to_orig to match normalized string exactly
     std::vector<size_t> norm_to_orig;
     norm_to_orig.reserve(input.size() * 3);
 
@@ -550,10 +553,8 @@ struct SpmUgmTokenizer {
     bool is_space_prepended = false;
     bool processing_non_ws = false;
 
-    size_t input_len = input.size();
-
     std::string_view input_view(input);
-    size_t orig_offset = 0;
+    size_t orig_pos = 0;  // Current position in original input
 
     while (!input_view.empty()) {
       auto p = case_encoder_->NormalizePrefix(input_view);
@@ -562,7 +563,7 @@ struct SpmUgmTokenizer {
       if (p.second == 0 && p.first.empty()) {
         // Advance by one UTF-8 character to prevent infinite loop
         size_t skip = std::min(ustring::UTF8Len(input_view[0]), input_view.size());
-        orig_offset += skip;
+        orig_pos += skip;
         input_view.remove_prefix(skip);
         continue;
       }
@@ -574,39 +575,41 @@ struct SpmUgmTokenizer {
             processing_non_ws = true;
             if ((shall_prepend_space && !is_space_prepended) || shall_merge_spaces) {
               normalized.append(space);
-              // Track original position for space characters
               for (size_t j = 0; j < space.size(); j++) {
-                norm_to_orig.push_back(orig_offset);
+                norm_to_orig.push_back(orig_pos);
               }
               is_space_prepended = true;
             }
           }
           normalized.push_back(c);
-          norm_to_orig.push_back(orig_offset + i);
+          norm_to_orig.push_back(orig_pos);
         } else {
           if (processing_non_ws) {
             processing_non_ws = false;
           }
           if (!shall_merge_spaces) {
             normalized.append(space);
             for (size_t j = 0; j < space.size(); j++) {
-              norm_to_orig.push_back(orig_offset + i);
+              norm_to_orig.push_back(orig_pos);
             }
           }
         }
       }
 
-      orig_offset += static_cast<size_t>(p.second);
+      orig_pos += static_cast<size_t>(p.second);
       input_view.remove_prefix(static_cast<size_t>(p.second));
     }
 
+    // Ensure norm_to_orig matches normalized size before PostProcess
+    // This is critical for PostProcess to work correctly
+    while (norm_to_orig.size() < normalized.size()) {
+      norm_to_orig.push_back(orig_pos > 0 ? orig_pos - 1 : 0);
+    }
+
     case_encoder_->PostProcess(&normalized, &norm_to_orig);
 
     if (shall_append_space) {
       normalized.append(space);
-      for (size_t j = 0; j < space.size(); j++) {
-        norm_to_orig.push_back(input.size());
-      }
     }
 
     return normalized;