<commit_before>module Predictor::Base
  def self.included(base)
    base.extend(ClassMethods)
  end

  module ClassMethods
    def input_matrix(key, opts={})
      @matrices ||= {}
      @matrices[key] = opts
    end

    def limit_similarities_to(val)
      @similarity_limit_set = true
      @similarity_limit     = val
    end

    def similarity_limit
      @similarity_limit_set ? @similarity_limit : 128
    end

    def reset_similarity_limit!
      @similarity_limit_set = nil
      @similarity_limit     = nil
    end

    def input_matrices=(val)
      @matrices = val
    end

    def input_matrices
      @matrices
    end
  end

  def input_matrices
    @input_matrices ||= Hash[self.class.input_matrices.map{ |key, opts|
      opts.merge!(:key => key, :redis_prefix => redis_prefix)
      [ key, Predictor::InputMatrix.new(opts) ]
    }]
  end

  def redis_prefix
    "predictor"
  end

  def similarity_limit
    self.class.similarity_limit
  end

  def redis_key(*append)
    ([redis_prefix] + append).flatten.compact.join(":")
  end

  def method_missing(method, *args)
    if input_matrices.has_key?(method)
      input_matrices[method]
    else
      raise NoMethodError.new(method.to_s)
    end
  end

  def respond_to?(method)
    input_matrices.has_key?(method) ? true : super
  end

  def all_items
    Predictor.redis.smembers(redis_key(:all_items))
  end

  def add_to_matrix(matrix, set, *items)
    items = items.flatten if items.count == 1 && items[0].is_a?(Array)  # Old syntax
    input_matrices[matrix].add_to_set(set, *items)
  end

  def add_to_matrix!(matrix, set, *items)
    items = items.flatten if items.count == 1 && items[0].is_a?(Array)  # Old syntax
    add_to_matrix(matrix, set, *items)
    process_items!(*items)
  end

  def related_items(item)
    keys = []
    input_matrices.each do |key, matrix|
      sets = Predictor.redis.smembers(matrix.redis_key(:sets, item))
      keys.concat(sets.map { |set| matrix.redis_key(:items, set) })
    end

    keys.empty? ? [] : (Predictor.redis.sunion(keys) - [item.to_s])
  end

  def predictions_for(set=nil, item_set: nil, matrix_label: nil, with_scores: false, offset: 0, limit: -1, exclusion_set: [], boost: {})
    fail "item_set or matrix_label and set is required" unless item_set || (matrix_label && set)

    if matrix_label
      matrix = input_matrices[matrix_label]
      item_set = Predictor.redis.smembers(matrix.redis_key(:items, set))
    end

    item_keys = []
    weights   = []

    item_set.each do |item|
      item_keys << redis_key(:similarities, item)
      weights   << 1.0
    end

    return [] if item_keys.empty?

    boost.each do |matrix_label, values|
      m = input_matrices[matrix_label]

      case values
      when Hash
        values[:values].each do |value|
          item_keys << m.redis_key(:items, value)
          weights   << values[:weight]
        end
      when Array
        values.each do |value|
          item_keys << m.redis_key(:items, value)
          weights   << 1.0
        end
      else
        raise "Bad value for boost: #{boost.inspect}"
      end
    end

    predictions = nil

    Predictor.redis.multi do |multi|
      # Passing plain sets to zunionstore is undocumented, but tested and supported:
      # https://github.com/antirez/redis/blob/2.8.11/tests/unit/type/zset.tcl#L481-L489

      multi.zunionstore 'temp', item_keys<commit_msg>Move comment.<commit_after>module Predictor::Base
  def self.included(base)
    base.extend(ClassMethods)
  end

  module ClassMethods
    def input_matrix(key, opts={})
      @matrices ||= {}
      @matrices[key] = opts
    end

    def limit_similarities_to(val)
      @similarity_limit_set = true
      @similarity_limit     = val
    end

    def similarity_limit
      @similarity_limit_set ? @similarity_limit : 128
    end

    def reset_similarity_limit!
      @similarity_limit_set = nil
      @similarity_limit     = nil
    end

    def input_matrices=(val)
      @matrices = val
    end

    def input_matrices
      @matrices
    end
  end

  def input_matrices
    @input_matrices ||= Hash[self.class.input_matrices.map{ |key, opts|
      opts.merge!(:key => key, :redis_prefix => redis_prefix)
      [ key, Predictor::InputMatrix.new(opts) ]
    }]
  end

  def redis_prefix
    "predictor"
  end

  def similarity_limit
    self.class.similarity_limit
  end

  def redis_key(*append)
    ([redis_prefix] + append).flatten.compact.join(":")
  end

  def method_missing(method, *args)
    if input_matrices.has_key?(method)
      input_matrices[method]
    else
      raise NoMethodError.new(method.to_s)
    end
  end

  def respond_to?(method)
    input_matrices.has_key?(method) ? true : super
  end

  def all_items
    Predictor.redis.smembers(redis_key(:all_items))
  end

  def add_to_matrix(matrix, set, *items)
    items = items.flatten if items.count == 1 && items[0].is_a?(Array)  # Old syntax
    input_matrices[matrix].add_to_set(set, *items)
  end

  def add_to_matrix!(matrix, set, *items)
    items = items.flatten if items.count == 1 && items[0].is_a?(Array)  # Old syntax
    add_to_matrix(matrix, set, *items)
    process_items!(*items)
  end

  def related_items(item)
    keys = []
    input_matrices.each do |key, matrix|
      sets = Predictor.redis.smembers(matrix.redis_key(:sets, item))
      keys.concat(sets.map { |set| matrix.redis_key(:items, set) })
    end

    keys.empty? ? [] : (Predictor.redis.sunion(keys) - [item.to_s])
  end

  def predictions_for(set=nil, item_set: nil, matrix_label: nil, with_scores: false, offset: 0, limit: -1, exclusion_set: [], boost: {})
    fail "item_set or matrix_label and set is required" unless item_set || (matrix_label && set)

    if matrix_label
      matrix = input_matrices[matrix_label]
      item_set = Predictor.redis.smembers(matrix.redis_key(:items, set))
    end

    item_keys = []
    weights   = []

    item_set.each do |item|
      item_keys << redis_key(:similarities, item)
      weights   << 1.0
    end

    return [] if item_keys.empty?

    boost.each do |matrix_label, values|
      m = input_matrices[matrix_label]

      # Passing plain sets to zunionstore is undocumented, but tested and supported:
      # https://github.com/antirez/redis/blob/2.8.11/tests/unit/type/zset.tcl#L481-L489

      case values
      when Hash
        values[:values].each do |value|
          item_keys << m.redis_key(:items, value)
          weights   << values[:weight]
        end
      when Array
        values.each do |value|
          item_keys << m.redis_key(:items, value)
          weights   << 1.0
        end
      else
        raise "Bad value for boost: #{boost.inspect}"
      end
    end

    predictions = nil

    Predictor.redis.multi do |multi|
      multi.zunionstore 'temp', item_keys
{'fraction_non_alphanumeric': 0.11027679623085984, 'fraction_numerical': 0.006772673733804476, 'mean_word_length': 2.7928531546622, 'pattern_counts': {'":': 2, '<': 27, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6052877', 'n_tokens_mistral': 2412, 'n_tokens_neox': 2289, 'n_words': 500}