Class: Llmemory::Retrieval::MmrReranker

Inherits:
Object
  • Object
show all
Defined in:
lib/llmemory/retrieval/mmr_reranker.rb

Instance Method Summary collapse

Constructor Details

#initialize(lambda: 0.7) ⇒ MmrReranker

Returns a new instance of MmrReranker.



6
7
8
# File 'lib/llmemory/retrieval/mmr_reranker.rb', line 6

def initialize(lambda: 0.7)
  @lambda = lambda
end

Instance Method Details

#rerank(candidates, score_key: :temporal_score) ⇒ Object



10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# File 'lib/llmemory/retrieval/mmr_reranker.rb', line 10

def rerank(candidates, score_key: :temporal_score)
  return candidates if candidates.size <= 1

  selected = []
  remaining = candidates.dup

  while remaining.any?
    best_idx = nil
    best_mmr = -Float::INFINITY

    remaining.each_with_index do |cand, i|
      rel = (cand[score_key] || cand[score_key.to_s] || cand[:score] || cand["score"] || 0).to_f
      max_sim = selected.map { |s| similarity(cand, s) }.max || 0
      mmr = @lambda * rel - (1 - @lambda) * max_sim

      if mmr > best_mmr
        best_mmr = mmr
        best_idx = i
      end
    end

    break unless best_idx

    selected << remaining.delete_at(best_idx)
  end

  selected
end