wandb 0.1.2 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fc087b1331eb3548a0c0826543a956c44ac2f301cf67de163a8b98cfbd425dd1
-  data.tar.gz: 7320addea0abaa3851ae9a1990c078f78fdb3f55d284ab7b7617cbb4e2203d92
+  metadata.gz: 4a00bf9591759e1339bca2ff33aca8dc245f71e8ef8f9e3aa4cd70a59f598a10
+  data.tar.gz: 880bed08a72bb3f03d48d093ab68a94ab8b706b759ec10945008979fd0c044c8
 SHA512:
-  metadata.gz: 8225c6929a867bc6d7e3da95698fa4d71f7f87bf2e8b1004601aa9cb7c8214c327570531a292b5f5df5a85f608fb498944ef22093803f7e527c49f87af640f74
-  data.tar.gz: 22e9685e8b57ea486f19fc91d3253ebed100c708f40aa739155b7eb5ec6c664a3ef6108ead7c933cd5c6e0a3324cb91b1020f3765d7b44143d87ca5c80d72f4e
+  metadata.gz: b6148f71160761a123696f0beac8b2ebe68064a4f64bd55ce1e9ec417c44b5c032a9f60b77394a3ebd0f0535e60c2ce50e53b5606f8acb958bfd35148c0f4a53
+  data.tar.gz: 14a72f4b7b0bc05234d9e0e276c724b1077ac61c07b755f90ce365d4d596c9183928b5df6d9396391a9339d6441b92e1e4c0733f67ee19634ba7fec47acb6818

data/lib/wandb/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Wandb
-  VERSION = "0.1.2"
+  VERSION = "0.1.5"
 end

data/lib/wandb/xgboost_callback.rb

@@ -1,26 +1,54 @@
+require "xgb"
+require "tempfile"
+require "fileutils"
+
 module Wandb
-  class XGBoostCallback
+  class XGBoostCallback < XGBoost::TrainingCallback
     MINIMIZE_METRICS = %w[rmse logloss error] # Add other metrics as needed
     MAXIMIZE_METRICS = %w[auc accuracy] # Add other metrics as needed
 
-    def initialize(log_model: false, log_feature_importance: true, importance_type: "gain", define_metric: true)
-      @log_model = log_model
-      @log_feature_importance = log_feature_importance
-      @importance_type = importance_type
-      @define_metric = define_metric
+    class Opts
+      attr_accessor :options
+
+      def initialize(options = {})
+        @options = options
+      end
+
+      def default(key, default)
+        options.key?(key) ? options[key] : default
+      end
+    end
+
+    attr_accessor :project_name, :api_key
+
+    def initialize(options = {})
+      options = Opts.new(options)
+      @log_model = options.default(:log_model, false)
+      @log_feature_importance = options.default(:log_feature_importance, true)
+      @importance_type = options.default(:importance_type, "gain")
+      @define_metric = options.default(:define_metric, true)
+      @api_key = options.default(:api_key, ENV["WANDB_API_KEY"])
+      @project_name = options.default(:project_name, nil)
+
+      raise "WANDB_API_KEY required" unless api_key
+      raise "project_name required" unless project_name
+
+      Wandb.login(api_key: api_key)
+      Wandb.init(project: project_name)
 
       return if Wandb.current_run
 
       raise "You must call wandb.init() before WandbCallback()"
     end
 
-    def before_training(model:)
+    def before_training(model)
       # Update Wandb config with model configuration
       Wandb.current_run.config = model.params
       Wandb.log(model.params)
+      model
     end
 
-    def after_training(model:)
+    def after_training(model)
       # Log the model as an artifact
       log_model_as_artifact(model) if @log_model
 
@@ -28,54 +56,61 @@ module Wandb
       log_feature_importance(model) if @log_feature_importance
 
       # Log best score and best iteration
-      return unless model.best_score
+      unless model.best_score
+        Wandb.finish
+        return model
+      end
 
       Wandb.log(
         "best_score" => model.best_score.to_f,
         "best_iteration" => model.best_iteration.to_i
       )
+
+      Wandb.finish
+      FileUtils.rm_rf(Rails.root.join("wandb"))
+
+      model
     end
 
-    def before_iteration(model:, epoch:, evals:)
-      # noop
+    def before_iteration(_model, _epoch, _history)
+      false
     end
 
-    def after_iteration(model:, epoch:, evals:, res:)
-      res.each do |metric_name, value|
-        data, metric = metric_name.split("-", 2)
-        full_metric_name = "#{data}-#{metric}"
-
-        if @define_metric
-          define_metric(data, metric)
-          Wandb.log({ full_metric_name => value })
-        else
-          Wandb.log({ full_metric_name => value })
-        end
-      end
+    def after_iteration(_model, epoch, history)
+      history.each do |split, metric_scores|
+        metric = metric_scores.keys.first
+        values = metric_scores.values.last
+        epoch_value = values[epoch]
 
-      Wandb.log({ "epoch" => epoch })
-      @define_metric = false
+        define_metric(split, metric) if @define_metric && epoch == 0
+        full_metric_name = "#{split}-#{metric}"
+        Wandb.log({ full_metric_name => epoch_value })
+      end
+      Wandb.log("epoch" => epoch)
+      false
     end
 
     private
 
     def log_model_as_artifact(model)
-      model_name = "#{Wandb.current_run.id}_model.json"
-      model_path = File.join(Wandb.current_run.dir, model_name)
-      model.save_model(model_path)
-
-      model_artifact = Wandb.Artifact(name: model_name, type: "model")
-      model_artifact.add_file(model_path)
-      Wandb.current_run.log_artifact(model_artifact)
+      Dir.mktmpdir("wandb_xgboost_model") do |tmp_dir|
+        model_name = "model.json"
+        model_path = File.join(tmp_dir, model_name)
+        model.save_model(model_path)
+
+        model_artifact = Wandb.artifact(name: model_name, type: "model")
+        model_artifact.add_file(model_path)
+        Wandb.current_run.log_artifact(model_artifact)
+      end
     end
 
     def log_feature_importance(model)
       fi = model.score(importance_type: @importance_type)
       fi_data = fi.map { |k, v| [k, v] }
 
-      table = Wandb.Table(data: fi_data, columns: %w[Feature Importance])
-      bar_plot = Wandb.plot.bar(table, "Feature", "Importance", title: "Feature Importance")
-      Wandb.log({ "Feature Importance" => bar_plot })
+      table = Wandb::Table.new(data: fi_data, columns: %w[Feature Importance])
+      bar_plot = Wandb::Plot.bar(table.table, "Feature", "Importance", title: "Feature Importance")
+      Wandb.log({ "Feature Importance" => bar_plot.__pyptr__ })
     end
 
     def define_metric(data, metric_name)

data/lib/wandb.rb

@@ -16,10 +16,6 @@ module Wandb
       @wandb ||= PyCall.import_module("wandb")
     end
 
-    def Table(*args, **kwargs)
-      __pyptr__.Table.new(*args, **kwargs)
-    end
-
     def plot(*args, **kwargs)
       __pyptr__.plot(*args, **kwargs)
     end
@@ -30,11 +26,12 @@ module Wandb
     end
 
     # Expose wandb.Artifact
-    def Artifact(*args, **kwargs)
-      __pyptr__.Artifact.new(*args, **kwargs)
+    def artifact(*args, **kwargs)
+      py_artifact = __pyptr__.Artifact.new(*args, **kwargs)
+      Artifact.new(py_artifact)
     end
 
-    def Error
+    def error
       __pyptr__.Error
     end
 
@@ -72,6 +69,10 @@ module Wandb
     def api
       @api ||= Api.new(__pyptr__.Api.new)
     end
+
+    def plot
+      Plot
+    end
   end
 
   # Run class
@@ -80,6 +81,10 @@ module Wandb
       @run = run
     end
 
+    def run_id
+      @run.run_id
+    end
+
     def log(metrics = {})
       metrics.symbolize_keys!
       @run.log(metrics, {})
@@ -104,6 +109,53 @@ module Wandb
     def config=(new_config)
       @run.config.update(PyCall::Dict.new(new_config))
     end
+
+    def log_artifact(artifact)
+      @run.log_artifact(artifact.__pyptr__)
+    end
+  end
+
+  # Artifact class
+  class Artifact
+    def initialize(artifact)
+      @artifact = artifact
+    end
+
+    def __pyptr__
+      @artifact
+    end
+
+    def name
+      @artifact.name
+    end
+
+    def type
+      @artifact.type
+    end
+
+    def add_file(local_path, name = nil)
+      @artifact.add_file(local_path, name)
+    end
+
+    def add_dir(local_dir, name = nil)
+      @artifact.add_dir(local_dir, name)
+    end
+
+    def get_path(name)
+      @artifact.get_path(name)
+    end
+
+    def metadata
+      @artifact.metadata
+    end
+
+    def metadata=(new_metadata)
+      @artifact.metadata = new_metadata
+    end
+
+    def save
+      @artifact.save
+    end
   end
 
   # Api class
@@ -137,6 +189,78 @@ module Wandb
       @project.description
     end
   end
+
+  # Table class
+  class Table
+    attr_accessor :table, :data, :columns
+
+    def initialize(data: {}, columns: [])
+      @table = Wandb.__pyptr__.Table.new(data: data, columns: columns)
+      @data = data
+      @columns = columns
+    end
+
+    def __pyptr__
+      @table
+    end
+
+    def add_data(*args)
+      @table.add_data(*args)
+    end
+
+    def add_column(name, data)
+      @table.add_column(name, data)
+    end
+
+    def get_column(name)
+      @table.get_column(name)
+    end
+
+    def columns
+      @table.columns
+    end
+
+    def data
+      @table.data
+    end
+
+    def to_pandas
+      @table.get_dataframe
+    end
+  end
+
+  # Plot class
+  class Plot
+    class << self
+      def bar(table, x_key, y_key, title: nil)
+        py_plot = Wandb.__pyptr__.plot.bar(table.__pyptr__, x_key, y_key, title: title)
+        new(py_plot)
+      end
+
+      def line(table, x_key, y_key, title: nil)
+        py_plot = Wandb.__pyptr__.plot.line(table.__pyptr__, x_key, y_key, title: title)
+        new(py_plot)
+      end
+
+      def scatter(table, x_key, y_key, title: nil)
+        py_plot = Wandb.__pyptr__.plot.scatter(table.__pyptr__, x_key, y_key, title: title)
+        new(py_plot)
+      end
+
+      def histogram(table, value_key, title: nil)
+        py_plot = Wandb.__pyptr__.plot.histogram(table.__pyptr__, value_key, title: title)
+        new(py_plot)
+      end
+    end
+
+    def initialize(plot)
+      @plot = plot
+    end
+
+    def __pyptr__
+      @plot
+    end
+  end
 end
 
 require_relative "wandb/xgboost_callback"

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: wandb
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.5
 platform: ruby
 authors:
 - Brett Shollenberger
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-10-11 00:00:00.000000000 Z
+date: 2024-10-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pycall