Class: Secryst::Byt5Onnx

Inherits:
Object
  • Object
show all
Defined in:
lib/secryst/byt5_onnx.rb

Overview

Byte-level seq2seq (ByT5 family) inference over ONNX sessions from an IMF v1 model zip: sha256-verified graphs, greedy KV-cache decode when the zip ships decoder-kv.onnx (default), plain full-recompute fallback otherwise. Tokenization is the canonical ByT5 table (byte b -> id b+3, trailing EOS) — see Secryst::IMF.

Instance Method Summary collapse

Constructor Details

#initialize(zip_path) ⇒ Byt5Onnx

Returns a new instance of Byt5Onnx.



12
13
14
15
16
17
18
19
20
 
# File 'lib/secryst/byt5_onnx.rb', line 12

def initialize(zip_path)
  @manifest = IMF.manifest(zip_path)
  graphs = IMF.verify_and_read(zip_path)
  @tmpdir = Dir.mktmpdir('secryst-imf')
  @encoder = build_session(graphs, 'encoder.onnx')
  @kv = @manifest['decoder'] == 'kv' && graphs.key?('decoder-kv.onnx')
  @decoder = build_session(graphs, @kv ? 'decoder-kv.onnx' : 'decoder.onnx')
  @pasts = @kv ? zero_pasts : {}
end

Instance Method Details

#id ⇒ Object 



31
32
33
 
# File 'lib/secryst/byt5_onnx.rb', line 31

def id
  @manifest['id']
end

#translate(text, max_seq_length: 256) ⇒ Object 



22
23
24
25
26
27
28
29
 
# File 'lib/secryst/byt5_onnx.rb', line 22

def translate(text, max_seq_length: 256)
  ids = IMF.encode(text)
  return '' if ids.length == 1 # only the trailing EOS: empty input

  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)
end