diff --git a/lib/interscript/ml/byt5_onnx.rb b/lib/interscript/ml/byt5_onnx.rb index 63fe0ad4..6a9f837e 100644 --- a/lib/interscript/ml/byt5_onnx.rb +++ b/lib/interscript/ml/byt5_onnx.rb @@ -24,7 +24,7 @@ def translate(text, max_seq_length: 256) hidden = @encoder.predict({ input_ids: [ids] })['last_hidden_state'] tokens = @kv ? greedy_kv(hidden, max_seq_length) : greedy_plain(hidden, max_seq_length) - IMF.decode(tokens) + IMF.first_alternates(IMF.decode(tokens)) end def id diff --git a/lib/interscript/ml/imf.rb b/lib/interscript/ml/imf.rb index 3eea999d..6f120a8f 100644 --- a/lib/interscript/ml/imf.rb +++ b/lib/interscript/ml/imf.rb @@ -28,6 +28,13 @@ def encode(text) text.bytes.map { |b| b + BYTE_OFFSET } + [EOS_ID] end + # Seq2seq output contract: models trained on multi-reference + # corpora emit '/'-separated alternates per token; the runtime + # returns the primary choice only. + def first_alternates(text) + text.split.map { |tok| tok.split('/', 2).first }.join(' ') + end + def decode(token_ids) out = +"" token_ids.each do |token| diff --git a/spec/first_alternates_spec.rb b/spec/first_alternates_spec.rb new file mode 100644 index 00000000..75854544 --- /dev/null +++ b/spec/first_alternates_spec.rb @@ -0,0 +1,19 @@ +require "spec_helper" +require "interscript/ml/imf" + +RSpec.describe Interscript::ML::IMF do + describe ".first_alternates" do + it "keeps the primary choice per token" do + expect(described_class.first_alternates("وِكَالَةُ/وِكَالَةُ اَلْأَرْصَادِ")).to eq("وِكَالَةُ اَلْأَرْصَادِ") + end + + it "leaves plain text untouched" do + expect(described_class.first_alternates("مرحبا بالعالم")).to eq("مرحبا بالعالم") + end + + it "handles edges" do + expect(described_class.first_alternates("")).to eq("") + expect(described_class.first_alternates("a/b c/d")).to eq("a c") + end + end +end