reranker-ruby 0.1.0 → 0.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0b72a0d8a77aa8b43a90df33ef1aca6e74f886795cedb9decbc2f412f6d788da
-  data.tar.gz: 236483d06f4ca930556ef2bf90a757a7b171cb5175a479cf538086030dbd71a0
+  metadata.gz: bcd9ac5bc1d098f8780a6523e648505dbf0259f4335b07e909898203b9723866
+  data.tar.gz: 12297c4398121fe27b4f0ed148268197f44258dde3b9e5f7efcf305202a5e861
 SHA512:
-  metadata.gz: 5bed85be6192b79df97d15732b17d688d596393099e0b8ef70d117f686151fef88a409254726fdd090af0f1f708e208b070ae55867fb2d8c7ca92f328a75b0be
-  data.tar.gz: a1ed43eba939509a20439604d0783ebd49942c6b01a4934ca576d91d29aab156a4703b7ebafd00812386f0f80cb4d967c7e8184f429220074d7a8f52241f6d01
+  metadata.gz: 75a6a774a12b00ed332f5f401451a53164f8ef3d8fe23428d70cfc224fa951bf0fd7a655f72178a98194e0578698e9f1de826afcd09b59b2dee2ec2dd0ea494f
+  data.tar.gz: c46089ed7e3ebedc6aea3bd2455bac70cd3b1666b77519a970510a4d4790d2bb3857a5d295d8c35e88358234d39bee5fee8e276fa5b34d8511b47b47a10f87fb

data/README.md

@@ -1,18 +1,32 @@
 # reranker-ruby
 
-Cross-encoder reranking for Ruby RAG pipelines. The single biggest quality improvement you can add to retrieval-augmented generation.
+Cross-encoder reranking for Ruby RAG pipelines.
 
-Vector search finds candidates fast (approximate). Reranking makes them accurate (precise).
+After vector search retrieves candidate documents, a reranker scores each candidate against the query using a cross-encoder model, producing far more accurate relevance rankings than embedding similarity alone. This is the single biggest quality improvement you can add to a RAG pipeline.
+
+```
+Bi-encoder (embedding search):   score = cosine(embed(query), embed(doc))  — fast, approximate
+Cross-encoder (reranking):       score = model(query + doc)                 — slow, precise
+```
+
+The pattern: use bi-encoder for top-100 retrieval, then cross-encoder to rerank to top-10.
 
 ## Installation
 
+Add to your Gemfile:
+
 ```ruby
 gem "reranker-ruby"
 ```
 
-## Usage
+For local ONNX inference, also install:
 
-### Cohere Rerank
+```ruby
+gem "onnxruntime"
+gem "tokenizers"
+```
+
+## Quick Start
 
 ```ruby
 require "reranker_ruby"
@@ -33,8 +47,22 @@ results = reranker.rerank(query, documents, top_k: 3)
 results.each do |r|
   puts "#{r.score.round(4)} | #{r.text[0..60]}"
 end
+# 0.9987 | Paris is the capital and largest city of France.
+# 0.8234 | The Eiffel Tower is located in Paris.
+# 0.6123 | Lyon is the second-largest city in France.
 ```
 
+## Providers
+
+### Cohere Rerank
+
+```ruby
+reranker = RerankerRuby::Cohere.new(api_key: ENV["COHERE_API_KEY"])
+results = reranker.rerank(query, documents, top_k: 3)
+```
+
+Uses [Cohere Rerank API v2](https://docs.cohere.com/reference/rerank) with the `rerank-v3.5` model by default.
+
 ### Jina Rerank
 
 ```ruby
@@ -42,7 +70,52 @@ reranker = RerankerRuby::Jina.new(api_key: ENV["JINA_API_KEY"])
 results = reranker.rerank(query, documents, top_k: 3)
 ```
 
-### Hash Documents with Metadata
+Uses `jina-reranker-v2-base-multilingual` by default.
+
+### Local ONNX Inference
+
+Run cross-encoder models locally without API calls. Models are auto-downloaded from HuggingFace Hub.
+
+```ruby
+reranker = RerankerRuby::Onnx.new(
+  model: "cross-encoder/ms-marco-MiniLM-L-6-v2"
+)
+results = reranker.rerank(query, documents, top_k: 3)
+```
+
+Or use a local model file:
+
+```ruby
+reranker = RerankerRuby::Onnx.new(
+  model_path: "/path/to/reranker.onnx",
+  tokenizer: "cross-encoder/ms-marco-MiniLM-L-6-v2"
+)
+```
+
+Supported models:
+- `cross-encoder/ms-marco-MiniLM-L-6-v2`
+- `cross-encoder/ms-marco-MiniLM-L-12-v2`
+- `BAAI/bge-reranker-base`
+- `BAAI/bge-reranker-large`
+- `BAAI/bge-reranker-v2-m3`
+
+Requires the `onnxruntime` and `tokenizers` gems.
+
+## Result Object
+
+Every reranker returns an array of `Result` objects, sorted by relevance (highest first):
+
+```ruby
+result.text      # => "Paris is the capital..."
+result.score     # => 0.9987
+result.index     # => 0 (position in the original document array)
+result.metadata  # => {} (preserved from input)
+result.to_h      # => { text: "...", score: 0.9987, index: 0, metadata: {} }
+```
+
+## Structured Documents with Metadata
+
+Pass hashes instead of strings. Metadata is preserved through reranking:
 
 ```ruby
 documents = [
@@ -54,38 +127,192 @@ results = reranker.rerank(query, documents, top_k: 3)
 results.first.metadata  # => { source: "wiki", id: "doc1" }
 ```
 
-### Reciprocal Rank Fusion
+## Reciprocal Rank Fusion
+
+Combine results from multiple retrieval strategies before reranking:
+
+```ruby
+vector_results = collection.search(embedding, top_k: 50)
+keyword_results = Article.where("content LIKE ?", "%#{query}%").limit(50)
+
+fused = RerankerRuby::RRF.fuse(
+  vector_results.map(&:id),
+  keyword_results.map(&:id),
+  k: 60
+)
+# => ranked array of IDs by combined relevance
+
+# Then rerank the fused results for final precision
+top_docs = fused.first(20).map { |id| Document.find(id) }
+final = reranker.rerank(query, top_docs.map(&:content), top_k: 5)
+```
+
+## Ensemble Reranking
+
+Combine multiple rerankers with weighted score aggregation:
+
+```ruby
+cohere = RerankerRuby::Cohere.new(api_key: ENV["COHERE_API_KEY"])
+jina = RerankerRuby::Jina.new(api_key: ENV["JINA_API_KEY"])
+
+ensemble = RerankerRuby::Ensemble.new(
+  rerankers: [cohere, jina],
+  weights: [0.6, 0.4],
+  normalize: :min_max  # :min_max, :softmax, :sigmoid, or :none
+)
+
+results = ensemble.rerank(query, documents, top_k: 5)
+```
+
+## Score Normalization
 
-Combine results from multiple retrieval strategies:
+Different models produce scores on different scales. Normalize them for comparison:
 
 ```ruby
-vector_results = ["doc1", "doc3", "doc5"]
-keyword_results = ["doc2", "doc1", "doc4"]
+results = reranker.rerank(query, documents)
+
+# Min-max to [0, 1]
+normalized = RerankerRuby::ScoreNormalizer.min_max(results)
 
-fused = RerankerRuby::RRF.fuse(vector_results, keyword_results, k: 60)
+# Softmax (scores sum to 1.0)
+normalized = RerankerRuby::ScoreNormalizer.softmax(results)
+
+# Sigmoid (each score independently mapped to [0, 1])
+normalized = RerankerRuby::ScoreNormalizer.sigmoid(results)
 ```
 
-### Caching
+## Batch Reranking
+
+Rerank multiple queries concurrently:
 
 ```ruby
+queries = ["capital of France?", "tallest building?", "largest ocean?"]
+
+results = RerankerRuby::Batch.rerank(
+  reranker, queries, documents,
+  top_k: 5,
+  threads: 4
+)
+
+results[0]  # => results for queries[0]
+results[1]  # => results for queries[1]
+```
+
+## Caching
+
+Avoid duplicate API calls for the same query+documents:
+
+```ruby
+# In-memory cache
 reranker = RerankerRuby::Cohere.new(
   api_key: ENV["COHERE_API_KEY"],
   cache: RerankerRuby::Cache::Memory.new(ttl: 3600)
 )
 
+# Redis cache
+require "redis"
+reranker = RerankerRuby::Cohere.new(
+  api_key: ENV["COHERE_API_KEY"],
+  cache: RerankerRuby::Cache::Redis.new(redis: Redis.new, ttl: 3600)
+)
+
 reranker.rerank(query, docs, top_k: 5)  # API call
 reranker.rerank(query, docs, top_k: 5)  # cache hit
 ```
 
-## Result Object
+## Logging & Metrics
+
+Every rerank call is automatically instrumented:
 
 ```ruby
-result.text      # => "Paris is the capital..."
-result.score     # => 0.9987
-result.index     # => 0 (original position)
-result.metadata  # => {} (preserved from input)
+# Set log level
+RerankerRuby::Logging.logger = Logger.new($stdout)
+RerankerRuby::Logging.logger.level = Logger::INFO
+
+# Subscribe to rerank events
+RerankerRuby::Logging.on_rerank do |event|
+  puts "#{event[:reranker]} reranked #{event[:document_count]} docs in #{event[:duration_ms]}ms"
+  # event keys: :reranker, :query, :document_count, :top_k,
+  #             :result_count, :duration_ms, :top_score
+end
+```
+
+## Rails Integration
+
+### Configuration
+
+Run the install generator:
+
+```bash
+rails generate reranker_ruby:install
+```
+
+This creates `config/initializers/reranker_ruby.rb`:
+
+```ruby
+RerankerRuby.configure do |config|
+  config.default_provider = :cohere         # :cohere, :jina, or :onnx
+  config.cohere_api_key = ENV["COHERE_API_KEY"]
+  config.default_top_k = 10
+  config.cache_store = :memory              # :memory, :redis, or nil
+  config.cache_ttl = 3600
+end
+```
+
+Then use the global convenience method anywhere:
+
+```ruby
+results = RerankerRuby.rerank("What is Ruby?", documents, top_k: 5)
 ```
 
+### ActiveJob for Async Reranking
+
+For large result sets, run reranking in the background:
+
+```ruby
+RerankerRuby::RerankJob.perform_later(
+  query: "What is Ruby?",
+  documents: ["doc1", "doc2", ...],
+  top_k: 5,
+  callback: "MyRerankCallback"
+)
+
+# Callback class
+class MyRerankCallback
+  def self.on_rerank_complete(query, results)
+    # results is an array of hashes: [{ text:, score:, index:, metadata: }, ...]
+  end
+end
+```
+
+### Pipeline Middleware
+
+Plug into any RAG pipeline as a reranking step:
+
+```ruby
+middleware = RerankerRuby::Middleware.new(
+  reranker: RerankerRuby::Cohere.new(api_key: "..."),
+  top_k: 5,
+  text_key: :content
+)
+
+# Works with hashes, strings, or objects
+candidates = [
+  { content: "Paris is the capital...", source: "wiki" },
+  { content: "Berlin is the capital...", source: "wiki" },
+]
+
+results = middleware.call(query: "capital of France?", candidates: candidates)
+```
+
+## Dependencies
+
+**Runtime:** `net/http` (stdlib), `json` (stdlib), `logger`
+
+**Optional:** `onnxruntime` and `tokenizers` (for local ONNX inference), `redis` (for Redis caching)
+
+**Development:** `minitest`, `rake`, `webmock`
+
 ## License
 
 MIT

data/lib/reranker_ruby/base.rb

@@ -3,6 +3,7 @@
 require "net/http"
 require "json"
 require "uri"
+require "digest"
 
 module RerankerRuby
   class Error < StandardError; end
@@ -20,6 +21,12 @@ module RerankerRuby
 
     private
 
+    def validate_inputs!(query, documents, top_k)
+      raise ArgumentError, "query cannot be nil or empty" if query.nil? || query.to_s.strip.empty?
+      raise ArgumentError, "documents cannot be nil or empty" if documents.nil? || documents.empty?
+      raise ArgumentError, "top_k must be positive" if top_k && top_k <= 0
+    end
+
     def instrument(query:, document_count:, top_k:, &block)
       Logging.instrument(
         reranker_class: self.class.name,
@@ -40,15 +47,14 @@ module RerankerRuby
       document.reject { |k, _| k == :text || k == "text" }
     end
 
-    def cache_key(query, documents)
-      require "digest"
-      Digest::SHA256.hexdigest("#{query}:#{documents.map(&:to_s).join("|")}")
+    def cache_key(query, documents, top_k = nil)
+      Digest::SHA256.hexdigest("#{query}:#{top_k}:#{documents.map(&:to_s).join("|")}")
     end
 
-    def with_cache(query, documents, &block)
+    def with_cache(query, documents, top_k: nil, &block)
       return yield unless @cache
 
-      key = cache_key(query, documents)
+      key = cache_key(query, documents, top_k)
       cached = @cache.get(key)
       return cached if cached
 
@@ -59,21 +65,43 @@ module RerankerRuby
 
     def post(url, body, headers: {})
       uri = URI.parse(url)
-      http = Net::HTTP.new(uri.host, uri.port)
-      http.use_ssl = uri.scheme == "https"
-
-      request = Net::HTTP::Post.new(uri.path)
-      request["Content-Type"] = "application/json"
-      headers.each { |k, v| request[k] = v }
-      request.body = JSON.generate(body)
-
-      response = http.request(request)
-
-      unless response.is_a?(Net::HTTPSuccess)
-        raise APIError, "HTTP #{response.code}: #{response.body}"
+      retries = 0
+      max_retries = 3
+
+      begin
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = uri.scheme == "https"
+        http.open_timeout = 30
+        http.read_timeout = 30
+        http.write_timeout = 30
+
+        request = Net::HTTP::Post.new(uri.path)
+        request["Content-Type"] = "application/json"
+        request["User-Agent"] = "RerankerRuby/#{RerankerRuby::VERSION}"
+        headers.each { |k, v| request[k] = v }
+        request.body = JSON.generate(body)
+
+        response = http.request(request)
+
+        if response.code.to_i == 429 || response.code.to_i >= 500
+          raise APIError, "HTTP #{response.code}: #{response.body}"
+        end
+
+        unless response.is_a?(Net::HTTPSuccess)
+          raise APIError, "HTTP #{response.code}: #{response.body}"
+        end
+
+        JSON.parse(response.body)
+      rescue APIError => e
+        retries += 1
+        if retries <= max_retries && (e.message.include?("429") || e.message.include?("50"))
+          sleep(2 ** (retries - 1))
+          retry
+        end
+        raise
+      rescue JSON::ParserError => e
+        raise APIError, "Invalid JSON response: #{e.message}"
       end
-
-      JSON.parse(response.body)
     end
   end
 end

data/lib/reranker_ruby/batch.rb

@@ -21,6 +21,7 @@ module RerankerRuby
       end
 
       results = Array.new(queries.length)
+      errors = []
       mutex = Mutex.new
       queue = Queue.new
 
@@ -30,14 +31,19 @@ module RerankerRuby
       workers = threads.times.map do
         Thread.new do
           while (item = queue.pop)
-            query, idx = item
-            result = reranker.rerank(query, documents, top_k: top_k)
-            mutex.synchronize { results[idx] = result }
+            begin
+              query, idx = item
+              result = reranker.rerank(query, documents, top_k: top_k)
+              mutex.synchronize { results[idx] = result }
+            rescue => e
+              mutex.synchronize { errors << e }
+            end
           end
         end
       end
 
       workers.each(&:join)
+      raise errors.first if errors.any?
       results
     end
   end

data/lib/reranker_ruby/cache/memory.rb

@@ -6,30 +6,39 @@ module RerankerRuby
       def initialize(ttl: 3600)
         @ttl = ttl
         @store = {}
+        @mutex = Mutex.new
       end
 
       def get(key)
-        entry = @store[key]
-        return nil unless entry
+        @mutex.synchronize do
+          entry = @store[key]
+          return nil unless entry
 
-        if Time.now.to_f - entry[:time] > @ttl
-          @store.delete(key)
-          return nil
-        end
+          if Time.now.to_f - entry[:time] > @ttl
+            @store.delete(key)
+            return nil
+          end
 
-        entry[:value]
+          entry[:value]
+        end
       end
 
       def set(key, value)
-        @store[key] = { value: value, time: Time.now.to_f }
+        @mutex.synchronize do
+          @store[key] = { value: value, time: Time.now.to_f }
+        end
       end
 
       def clear
-        @store.clear
+        @mutex.synchronize do
+          @store.clear
+        end
       end
 
       def size
-        @store.size
+        @mutex.synchronize do
+          @store.size
+        end
       end
     end
   end

data/lib/reranker_ruby/cache/redis.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "json"
+
 module RerankerRuby
   module Cache
     class Redis
@@ -13,11 +15,20 @@ module RerankerRuby
         data = @redis.get("#{@prefix}#{key}")
         return nil unless data
 
-        Marshal.load(data) # rubocop:disable Security/MarshalLoad
+        parsed = JSON.parse(data)
+        parsed.map do |h|
+          Result.new(
+            text: h["text"],
+            score: h["score"],
+            index: h["index"],
+            metadata: h["metadata"] || {}
+          )
+        end
       end
 
       def set(key, value)
-        @redis.setex("#{@prefix}#{key}", @ttl, Marshal.dump(value))
+        serialized = JSON.generate(value.map(&:to_h))
+        @redis.setex("#{@prefix}#{key}", @ttl, serialized)
       end
 
       def clear

data/lib/reranker_ruby/cohere.rb

@@ -12,8 +12,9 @@ module RerankerRuby
     end
 
     def rerank(query, documents, top_k: 10, model: nil)
+      validate_inputs!(query, documents, top_k)
       instrument(query: query, document_count: documents.length, top_k: top_k) do
-        with_cache(query, documents) do
+        with_cache(query, documents, top_k: top_k) do
           texts = extract_texts(documents)
 
           response = post(API_URL, {

data/lib/reranker_ruby/configuration.rb

@@ -66,7 +66,10 @@ module RerankerRuby
 
   class << self
     def configuration
-      @configuration ||= Configuration.new
+      @config_mutex ||= Mutex.new
+      @config_mutex.synchronize do
+        @configuration ||= Configuration.new
+      end
     end
 
     def configure
@@ -74,13 +77,19 @@ module RerankerRuby
     end
 
     def reset_configuration!
-      @configuration = Configuration.new
-      @reranker = nil
+      @config_mutex ||= Mutex.new
+      @config_mutex.synchronize do
+        @configuration = Configuration.new
+        @reranker = nil
+      end
     end
 
     # Global reranker instance built from configuration
     def reranker
-      @reranker ||= configuration.build_reranker
+      @config_mutex ||= Mutex.new
+      @config_mutex.synchronize do
+        @reranker ||= configuration.build_reranker
+      end
     end
 
     # Convenience method for quick reranking

data/lib/reranker_ruby/ensemble.rb

@@ -27,7 +27,8 @@ module RerankerRuby
     end
 
     def rerank(query, documents, top_k: 10)
-      with_cache(query, documents) do
+      validate_inputs!(query, documents, top_k)
+      with_cache(query, documents, top_k: top_k) do
         texts = extract_texts(documents)
 
         # Collect and normalize results from each reranker

data/lib/reranker_ruby/jina.rb

@@ -12,8 +12,9 @@ module RerankerRuby
     end
 
     def rerank(query, documents, top_k: 10, model: nil)
+      validate_inputs!(query, documents, top_k)
       instrument(query: query, document_count: documents.length, top_k: top_k) do
-        with_cache(query, documents) do
+        with_cache(query, documents, top_k: top_k) do
           texts = extract_texts(documents)
 
           response = post(API_URL, {

data/lib/reranker_ruby/onnx.rb

@@ -26,7 +26,8 @@ module RerankerRuby
     end
 
     def rerank(query, documents, top_k: 10)
-      with_cache(query, documents) do
+      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) }

data/lib/reranker_ruby/score_normalizer.rb

@@ -9,6 +9,10 @@ module RerankerRuby
       return results if results.empty?
 
       scores = results.map(&:score)
+      if scores.any? { |s| s.nan? || s.infinite? }
+        return results.map { |r| with_score(r, 0.0) }
+      end
+
       min = scores.min
       max = scores.max
       range = max - min
@@ -23,6 +27,10 @@ module RerankerRuby
       return results if results.empty?
 
       scores = results.map(&:score)
+      if scores.any? { |s| s.nan? || s.infinite? }
+        return results.map { |r| with_score(r, 0.0) }
+      end
+
       max_score = scores.max
       exps = scores.map { |s| Math.exp(s - max_score) } # subtract max for numerical stability
       sum = exps.sum
@@ -34,6 +42,13 @@ module RerankerRuby
 
     # Sigmoid normalization — each score independently mapped to [0, 1]
     def self.sigmoid(results)
+      return results if results.empty?
+
+      scores = results.map(&:score)
+      if scores.any? { |s| s.nan? || s.infinite? }
+        return results.map { |r| with_score(r, 0.0) }
+      end
+
       results.map { |r| with_score(r, 1.0 / (1.0 + Math.exp(-r.score))) }
     end
 

data/lib/reranker_ruby/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RerankerRuby
-  VERSION = "0.1.0"
+  VERSION = "0.1.1"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: reranker-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.1
 platform: ruby
 authors:
 - Johannes Dwi Cahyo