Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion lib/interscript/ml/byt5_onnx.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions lib/interscript/ml/imf.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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|
Expand Down
19 changes: 19 additions & 0 deletions spec/first_alternates_spec.rb
Original file line number Diff line number Diff line change
@@ -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
Loading