Class: RerankerRuby::Onnx

Inherits:
Base
  • Object
show all
Defined in:
lib/reranker_ruby/onnx.rb

Constant Summary collapse

DEFAULT_MODEL = 
"cross-encoder/ms-marco-MiniLM-L-6-v2"
MAX_LENGTH = 
512

Instance Method Summary collapse

Constructor Details

#initialize(model: nil, model_path: nil, tokenizer: nil, cache_dir: nil, **options) ⇒ Onnx

Returns a new instance of Onnx.



8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
 
# File 'lib/reranker_ruby/onnx.rb', line 8

def initialize(model: nil, model_path: nil, tokenizer: nil, cache_dir: nil, **options)
  super(**options)
  require_dependencies!

  if model_path
    @model_path = model_path
    tokenizer_id = tokenizer || DEFAULT_MODEL
    @tokenizer = Tokenizers.from_pretrained(tokenizer_id)
  else
    repo_id = model || DEFAULT_MODEL
    downloader = ModelDownloader.new(cache_dir: cache_dir || ModelDownloader::DEFAULT_CACHE_DIR)
    paths = downloader.download(repo_id)
    @model_path = paths[:model_path]
    @tokenizer = Tokenizers.from_file(paths[:tokenizer_path])
  end

  @session = OnnxRuntime::InferenceSession.new(@model_path)
  @tokenizer.enable_truncation(MAX_LENGTH)
end

Instance Method Details

#rerank(query, documents, top_k: 10) ⇒ Object 



28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/reranker_ruby/onnx.rb', line 28

def rerank(query, documents, top_k: 10)
  validate_inputs!(query, documents, top_k)
  with_cache(query, documents, top_k: top_k) do
    texts = extract_texts(documents)

    scores = texts.map { |text| score_pair(query, text) }

    results = texts.each_with_index.map do |text, idx|
      Result.new(
        text: text,
        score: scores[idx],
        index: idx,
        metadata: extract_metadata(documents[idx])
      )
    end

    results.sort.first(top_k)
  end
end