Class: Gigatoken::PackedResult

Inherits:
Object
  • Object
show all
Includes:
Enumerable
Defined in:
lib/gigatoken/packed_result.rb

Overview

A packed batch-encode result: one IO::Buffer of u32 token ids (native byte order) for every document, plus lens, each document's token count. Returned by Tokenizer#encode_batch/#encode_files with packed: true — the Ruby analog of the numpy/awkward-array path the native ext's header comment says was dropped for lack of a Ruby equivalent. Documents are unpacked from the buffer on demand via #[]/#each.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(buffer, lens) ⇒ PackedResult

Returns a new instance of PackedResult.



15
16
17
18
19
# File 'lib/gigatoken/packed_result.rb', line 15

def initialize(buffer, lens)
  @buffer = buffer
  @lens = lens
  @offsets = lens.inject([0]) { |offsets, len| offsets << offsets.last + len }
end

Instance Attribute Details

#bufferObject (readonly)

Returns the value of attribute buffer.



13
14
15
# File 'lib/gigatoken/packed_result.rb', line 13

def buffer
  @buffer
end

#lensObject (readonly)

Returns the value of attribute lens.



13
14
15
# File 'lib/gigatoken/packed_result.rb', line 13

def lens
  @lens
end

Instance Method Details

#[](i) ⇒ Object

Array of token ids for document i, materialized on demand.



32
33
34
# File 'lib/gigatoken/packed_result.rb', line 32

def [](i)
  buffer.get_values(Array.new(lens[i], :u32), @offsets[i] * 4)
end

#eachObject



36
37
38
39
40
# File 'lib/gigatoken/packed_result.rb', line 36

def each
  return enum_for(:each) unless block_given?

  lens.each_index { |i| yield self[i] }
end

#sizeObject

Number of documents.



22
23
24
# File 'lib/gigatoken/packed_result.rb', line 22

def size
  lens.size
end

#to_aObject

A ragged Array of Arrays, one per document — the same shape encode_batch/encode_files return with packed: false.



44
45
46
# File 'lib/gigatoken/packed_result.rb', line 44

def to_a
  each.to_a
end

#token_countObject

Total token count across every document.



27
28
29
# File 'lib/gigatoken/packed_result.rb', line 27

def token_count
  lens.sum
end