bm25f 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: bc3b5ddbf5d62479a1a0afafccef1d377db2a149ac9db87f8d207b0c16b41550
+  data.tar.gz: 295be306f71a84ae399cffaf7c1faff3a768ad7e1c00095230144d564cc9af59
+SHA512:
+  metadata.gz: 1c06abf6c6d53a66e151378c610bf5ecfbe92ced01692bbe126d7bfda77bd85a2c2733934460052a6d8129e61881e067a1e35264f7acb197ff51f2336bb31e19
+  data.tar.gz: 85ad4ba68f49009bbe667ccf993d498671ced59f9aeff4394435778c53369037f7dee346b7111c3cc6679fc72b9c8e62798139d3dce7a56e098795844e59b85e

data/lib/bm25f.rb

@@ -0,0 +1,141 @@
+require 'treat'
+
+class BM25F
+  include Treat::Core::DSL
+
+  # Initializes a BM25F model.
+  #
+  # @param term_freq_weight [Float] Weight for term frequency.
+  # @param doc_length_weight [Float] Weight for document length.
+  def initialize(term_freq_weight: 1.33, doc_length_weight: 0.8)
+    @term_freq_weight = term_freq_weight
+    @doc_length_weight = doc_length_weight
+  end
+
+  # Fits the model to a set of documents.
+  #
+  # @param documents [Hash] The documents to fit the model to.
+  # @param field_weights [Hash] A specified weight for each key the documents.
+  def fit(documents, field_weights = {})
+    documents = preprocess_documents(documents)
+
+    # Set missing field_weights to 1
+    unique_keys = documents.flat_map(&:keys).uniq
+
+    unique_keys.each do |key|
+      field_weights[key] = 1 unless field_weights.key?(key)
+    end
+
+    @field_weights = field_weights
+    @documents = documents
+    @avg_doc_length = calculate_average_document_length(documents)
+    @doc_lengths = calculate_document_lengths(documents)
+    @total_docs = documents.length
+    @idf = calculate_idf
+  end
+
+  # Calculates the score of each document using the query.
+  #
+  # @param query [String] The query to score with.
+  # @return [Hash] A hash containing document IDs and their scores.
+  def score(query)
+    query_terms = preprocess_query(query)
+    scores = {}
+    (0...@total_docs).each do |doc_id|
+      scores[doc_id] = calculate_document_score(doc_id, query_terms)
+    end
+    scores
+  end
+
+  private
+
+  # Preprocesses documents by tokenizing and stemming them.
+  #
+  # @param documents [Hash] The documents to preprocess.
+  def preprocess_documents(documents)
+    documents.each do |k, v|
+      next unless v.instance_of? String
+
+      documents[k] = sentence(v).map(&:stem).join(' ')
+    end
+    documents
+  end
+
+  # Calculates the average document length.
+  #
+  # @param documents [Hash] The documents.
+  # @return [Float] The average document length.
+  def calculate_average_document_length(documents)
+    total_length = documents.sum { |doc| doc.values.map(&:length).sum }
+    total_length / documents.length.to_f
+  end
+
+  # Calculates the lengths of each field in a document.
+  #
+  # @param documents [Hash] The documents.
+  # @return [Hash] A hash of document lengths.
+  def calculate_document_lengths(documents)
+    doc_lengths = {}
+    documents.each_with_index do |doc, i|
+      doc_lengths[i] = doc.transform_values(&:length)
+    end
+    doc_lengths
+  end
+
+  # Calculates the IDF for each field.
+  #
+  # @return [Hash] A hash of IDF values for each field.
+  def calculate_idf
+    idf = {}
+    @field_weights.each_key do |field|
+      field_doc_count = @documents.count { |doc| !doc[field].empty? }
+      idf[field] = Math.log((@total_docs - field_doc_count + 0.5) / (field_doc_count + 0.5) + 1.0)
+    end
+    idf
+  end
+
+  # Preprocesses a query by tokenizing and stemming it.
+  #
+  # @param query [String] The query to preprocess.
+  # @return [Array<String>] An array of preprocessed query terms.
+  def preprocess_query(query)
+    sentence(query).tokenize.map(&:stem)
+  end
+
+  # Calculates the score of a document using an array of query terms.
+  #
+  # @param doc_id [Integer] The document ID.
+  # @param query_terms [Array<String>] The query terms.
+  # @return [Float] The document score.
+  def calculate_document_score(doc_id, query_terms)
+    doc_score = 0
+    @field_weights.each_key do |field|
+      query_terms.each do |term|
+        tf = field_term_frequency(field, term, doc_id)
+        idf = @idf[field]
+        field_length_norm = field_length_norm(field, doc_id)
+        doc_score += @field_weights[field] * ((tf * (@term_freq_weight + 1)) / (tf + @term_freq_weight * field_length_norm) * idf)
+      end
+    end
+    doc_score
+  end
+
+  # Calculates the term frequency in a field of a document.
+  #
+  # @param field [Symbol] The field name.
+  # @param term [String] The term to calculate frequency for.
+  # @param doc_id [Integer] The document ID.
+  # @return [Integer] The term frequency.
+  def field_term_frequency(field, term, doc_id)
+    @documents[doc_id][field].scan(term).count
+  end
+
+  # Calculates the field length normalization factor of a document.
+  #
+  # @param field [Symbol] The field name.
+  # @param doc_id [Integer] The document ID.
+  # @return [Float] The field length normalization factor.
+  def field_length_norm(field, doc_id)
+    1.0 - @doc_length_weight + @doc_length_weight * (@doc_lengths[doc_id][field] / @avg_doc_length)
+  end
+end

metadata

@@ -0,0 +1,60 @@
+--- !ruby/object:Gem::Specification
+name: bm25f
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- catflip
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2023-09-09 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: treat
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.1'
+description: A fast implementation of the BM25F ranking algorithm for information
+  retrieval systems, written in Ruby.
+email:
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- lib/bm25f.rb
+homepage: https://github.com/catflip/bm25f-ruby
+licenses:
+- AGPL-3.0
+metadata:
+  homepage_uri: https://github.com/catflip/bm25f-ruby
+  source_code_uri: https://github.com/catflip/bm25f-ruby
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.0.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.3.26
+signing_key:
+specification_version: 4
+summary: BM25F ranking function in Ruby.
+test_files: []