Class: StaticEmbeddings::Format::TrieBuilder

Inherits:
Object
  • Object
show all
Defined in:
lib/static_embeddings/format.rb

Defined Under Namespace

Classes: Node

Instance Method Summary collapse

Constructor Details

#initialize ⇒ TrieBuilder

Returns a new instance of TrieBuilder.



184
185
186
 
# File 'lib/static_embeddings/format.rb', line 184

def initialize
  @nodes = [Node.new(terminal: SLOT_EMPTY, children: {})]
end

Instance Method Details

#insert(bytes, id) ⇒ Object

Raises:

  • (ArgumentError) 


188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
 
# File 'lib/static_embeddings/format.rb', line 188

def insert(bytes, id)
  return if bytes.empty?

  node_index = 0
  bytes.each_byte do |byte|
    node = @nodes[node_index]
    child = node.children[byte]
    unless child
      child = @nodes.length
      node.children[byte] = child
      @nodes << Node.new(terminal: SLOT_EMPTY, children: {})
    end
    node_index = child
  end

  node = @nodes[node_index]
  raise ArgumentError, "duplicate trie key for token id #{id}" unless node.terminal == SLOT_EMPTY

  node.terminal = id
end

#pack ⇒ Object 



209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
 
# File 'lib/static_embeddings/format.rb', line 209

def pack
  edges = []
  node_records = @nodes.map do |node|
    start = edges.length
    node.children.sort_by { |byte, _| byte }.each do |byte, child|
      edges << [byte, child]
    end
    [start, node.children.length, node.terminal, 0]
  end

  packed = Format.binary_string(16 + node_records.length * 16 + edges.length * 8)
  packed << [node_records.length, edges.length, 0, 0].pack("V4")
  node_records.each { |record| packed << record.pack("V4") }
  edges.each { |edge| packed << edge.pack("V2") }
  packed
end