feat: plane runtime preserve_diacritics mode — user diacritics pinned in decode, byte-exact round-trip (#801)

Author: ronaldtse

lib/interscript/ml/plane.rb

@@ -21,6 +21,10 @@ def initialize(graph:, classes:, k_passes: 2)
         @n_classes = classes.length
         @mask_id = @n_classes
         @k_passes = k_passes
+        # a character is a diacritic iff it occurs inside some class the
+        # artifact knows — no per-language mark tables needed
+        @mark_chars = classes.flat_map { |c| c.chars }.reject(&:empty?).uniq.freeze
+        @class_to_id = classes.each_with_index.to_h.freeze
         # the onnxruntime gem loads from paths, not bytes (as byt5_onnx):
         # verified bytes go to a tmpfile owned by this instance
         @graph_tmpdir = Dir.mktmpdir('isx-plane')
@@ -29,17 +33,69 @@ def initialize(graph:, classes:, k_passes: 2)
         @sess = OnnxRuntime::Model.new(path)
       end
 
-      def translate(text)
-        ids = IMF.encode(text)[0..-2] # strip EOS: one token per byte
-        n = ids.length
-        plane = Array.new(n, @mask_id)
-        preds = nil
-        @k_passes.times do
-          logits = @sess.predict({ 'input_ids' => [ids], 'plane_ids' => [plane] })['class_logits'][0]
-          preds = logits.each_with_index.map { |row, i| row.index(row.max) }
-          plane = preds
+      # Split text into (skeleton, per-base classes). Marks FOLLOW their
+      # base letter; leading marks bind to a "\x00" anchor (render drops
+      # it). Cluster order is preserved as written — render-exactness
+      # beats normalization (the Hebrew canon lesson).
+      def split_planes(text)
+        skeleton = []
+        classes = []
+        current = []
+        close = lambda do
+          next if current.empty?
+
+          if skeleton.empty?
+            skeleton << "\x00"
+            classes << current.join
+          else
+            classes[-1] = classes[-1] + current.join
+          end
+          current = []
+        end
+        text.each_char do |ch|
+          if @mark_chars.include?(ch)
+            current << ch
+          else
+            close.call
+            skeleton << ch
+            classes << ''
+          end
+        end
+        close.call
+        [skeleton.join, classes]
+      end
+
+      def translate(text, preserve_diacritics: false)
+        return translate_plain(text) unless preserve_diacritics
+
+        skeleton, char_classes = split_planes(text)
+        tok_preds = decode_pinned(skeleton, char_classes)
+        pos = 0
+        final = char_classes.each_with_index.map do |user_cls, i|
+          ch = skeleton[i]
+          n = ch.bytesize
+          cls = if user_cls.empty?
+                  votes = tok_preds[pos, n] || []
+                  cid = votes.tally.max_by { |_, c| c }&.first || @n_classes
+                  cid < @n_classes ? @classes[cid] : ''
+                else
+                  user_cls # user diacritics round-trip byte-exactly
+                end
+          pos += n
+          cls
+        end
+        out = +""
+        skeleton.each_char.with_index do |ch, i|
+          out << (ch == "\x00" ? final[i] : ch + final[i])
         end
+        out
+      end
+
+      private
 
+      def translate_plain(text)
+        ids = IMF.encode(text)[0..-2] # strip EOS: one token per byte
+        preds = k_pass_decode(ids, Array.new(ids.length, @mask_id))
         pos = 0
         char_preds = text.each_char.map do |ch|
           nbytes = ch.bytesize
@@ -51,6 +107,32 @@ def translate(text)
         PlaneModel.render_plane(text, char_preds)
       end
 
+      def k_pass_decode(ids, plane0)
+        plane = plane0
+        preds = nil
+        @k_passes.times do
+          logits = @sess.predict({ 'input_ids' => [ids], 'plane_ids' => [plane] })['class_logits'][0]
+          preds = logits.each_with_index.map { |row, i| row.index(row.max) }
+          plane = preds
+        end
+        preds
+      end
+
+      # K-pass decode with user classes pinned: known classes are
+      # conditioned AND never overwritten; unknown non-empty clusters
+      # stay masked during decode but are restored at render.
+      def decode_pinned(skeleton, char_classes)
+        ids = IMF.encode(skeleton)[0..-2]
+        per_token_pin = []
+        skeleton.each_char.with_index do |ch, i|
+          cls = char_classes[i]
+          pid = @class_to_id.fetch(cls, @mask_id)
+          ch.bytesize.times { per_token_pin << pid }
+        end
+        plane0 = per_token_pin.map { |p| p == @mask_id ? @mask_id : p }
+        k_pass_decode(ids, plane0)
+      end
+
       def self.render_plane(skeleton, char_classes)
         skeleton.each_char.with_index.map do |ch, i|
           ch + char_classes[i].to_s

spec/interscript/ml/plane_spec.rb

@@ -50,3 +50,37 @@
     end
   end
 end
+
+RSpec.describe Interscript::ML::PlaneModel, "#translate preserve_diacritics" do
+  let(:zip) { File.expand_path("../fixtures/tiny-plane.zip", __dir__) }
+  let(:m) { described_class.from_zip(File.read(zip)) }
+  let(:plain_of) { ->(s) { s.chars.reject { |c| m.classes.include?(c) }.join } }
+
+  it "default path unchanged by flag" do
+    expect(m.translate("كتب")).to eq(m.translate("كتب", preserve_diacritics: false))
+  end
+
+  it "fully labeled input round-trips byte-exactly" do
+    labeled = "كُتُبُ"
+    expect(m.translate(labeled, preserve_diacritics: true)).to eq(labeled)
+  end
+
+  it "partial input keeps user classes and fills the rest" do
+    out = m.translate("كُتب", preserve_diacritics: true)
+    expect(out).to start_with("كُ")
+    expect(plain_of.call(out)).to eq("كتب")
+  end
+
+  it "leading marks round-trip without anchors leaking" do
+    out = m.translate("ُكتب", preserve_diacritics: true)
+    expect(out[0]).to eq("ُ")
+    expect(out).not_to include("\x00")
+    expect(plain_of.call(out)).to eq("كتب")
+  end
+
+  it "unknown cluster still round-trips byte-exactly" do
+    out = m.translate("كُْتب", preserve_diacritics: true)
+    expect(out).to include("كُْ")
+    expect(plain_of.call(out)).to eq("كتب")
+  end
+end