Microservice using Sinatra framework to embed text using Multilingual E5 Large embeddings model.

app.rb

require 'json'

require 'sinatra'

require 'onnxruntime'
require 'tokenizers'

require_relative './embedder.rb'

before do
  unless defined?(SS_MODEL) && defined?(SS_TOKENIZER)
    OnnxRuntime.ffi_lib = 'onnxruntime-linux-x64-gpu-1.17.0/lib/libonnxruntime.so'

    SS_MODEL = OnnxRuntime::Model.new("multilingual-e5-large-onnx/model.onnx", providers: ['CUDAExecutionProvider'])
    SS_TOKENIZER = Tokenizers.from_file("multilingual-e5-large-onnx/tokenizer.json")

    SS_TOKENIZER.enable_truncation(512)
  end
end

post '/embed' do
  body = JSON.parse(request.body.read)

  [200, Embedder.embed(body['text'], type: body['type']).to_json]
end

embedder.rb

class Embedder
  def self.embed(text, type:)
    token_ids = SS_TOKENIZER.encode("#{type}: #{text}").ids
    outputs = SS_MODEL.predict({ input_ids: [token_ids], attention_mask: [[1] * token_ids.size] })

    l2_normalize(average_pool(outputs['last_hidden_state'][0]))
  end

  def self.average_pool(last_hidden_states)
    last_hidden_states.transpose.map { |values| values.sum.to_f / values.size }
  end

  def self.l2_normalize(vector)
    norm = Math.sqrt(vector.map { |x| x**2 }.sum)
    vector.map { |x| x / norm }
  end
end