Method: Vecsearch::GTETiny#encode_batch

Defined in:
lib/vecsearch/gte_tiny.rb

#encode_batch(input, n_threads: 1) ⇒ Object



57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# File 'lib/vecsearch/gte_tiny.rb', line 57

def encode_batch(input, n_threads: 1)
  # Create an array of pointers to the input strings
  input_ptrs = input.map { |str| FFI::MemoryPointer.from_string(str) }

  # Create a pointer to the array of input pointers
  input_ptrs_ptr = FFI::MemoryPointer.new(:pointer, input_ptrs.length)
  input_ptrs_ptr.write_array_of_pointer(input_ptrs)

  # Create an output buffer for each input string
  output_ptrs = input.map { FFI::MemoryPointer.new(:float, @n_embd) }

  # Create a pointer to the array of output pointers
  output_ptrs_ptr = FFI::MemoryPointer.new(:pointer, output_ptrs.length)
  output_ptrs_ptr.write_array_of_pointer(output_ptrs)

  Bert.bert_encode_batch(@ctx, n_threads, MAX_TOKENS, input.length, input_ptrs_ptr, output_ptrs_ptr)

  # Convert the output buffers to Ruby arrays
  output = output_ptrs.map { |ptr| ptr.read_array_of_float(@n_embd) }

  output
end