├── .github └── workflows │ └── build.yml ├── .gitignore ├── CHANGELOG.md ├── Gemfile ├── LICENSE.txt ├── README.md ├── Rakefile ├── lib ├── torchrec.rb └── torchrec │ ├── models │ ├── deepfm │ │ ├── dense_arch.rb │ │ └── over_arch.rb │ └── dlrm │ │ └── dense_arch.rb │ ├── modules │ ├── activation │ │ └── swish_layer_norm.rb │ ├── cross_net │ │ └── cross_net.rb │ ├── deepfm │ │ ├── deepfm.rb │ │ └── factorization_machine.rb │ ├── mlp │ │ ├── mlp.rb │ │ └── perceptron.rb │ └── utils.rb │ ├── sparse │ └── jagged_tensor.rb │ └── version.rb ├── test ├── models_test.rb ├── modules_test.rb ├── sparse_test.rb └── test_helper.rb └── torchrec.gemspec /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | on: [push, pull_request] 3 | jobs: 4 | build: 5 | runs-on: ubuntu-latest 6 | env: 7 | BUNDLE_BUILD__TORCH___RB: "--with-torch-dir=/home/runner/libtorch" 8 | LIBTORCH_VERSION: 2.5.1 9 | steps: 10 | - uses: actions/checkout@v4 11 | - uses: actions/cache@v4 12 | with: 13 | path: ~/libtorch 14 | key: libtorch-${{ env.LIBTORCH_VERSION }} 15 | id: cache-libtorch 16 | - name: Download LibTorch 17 | if: steps.cache-libtorch.outputs.cache-hit != 'true' 18 | run: | 19 | cd ~ 20 | wget -q -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-$LIBTORCH_VERSION%2Bcpu.zip 21 | unzip -q libtorch.zip 22 | - uses: ruby/setup-ruby@v1 23 | with: 24 | ruby-version: 3.4 25 | bundler-cache: true 26 | - run: bundle exec rake test 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /.bundle/ 2 | /.yardoc 3 | /_yardoc/ 4 | /coverage/ 5 | /doc/ 6 | /pkg/ 7 | /spec/reports/ 8 | /tmp/ 9 | *.lock 10 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## 0.0.4 (2024-08-02) 2 | 3 | - Dropped support for Ruby < 3.1 4 | 5 | ## 0.0.3 (2023-07-24) 6 | 7 | - Dropped support for Ruby < 3 8 | 9 | ## 0.0.2 (2022-03-14) 10 | 11 | - Added `JaggedTensor` 12 | - Added `CrossNet`, `DeepFM`, and `FactorizationMachine` modules 13 | 14 | ## 0.0.1 (2022-02-28) 15 | 16 | - First release 17 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source "https://rubygems.org" 2 | 3 | gemspec 4 | 5 | gem "rake" 6 | gem "minitest" 7 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) Meta Platforms, Inc. and affiliates. 4 | Copyright (c) Andrew Kane. 5 | All rights reserved. 6 | 7 | Redistribution and use in source and binary forms, with or without 8 | modification, are permitted provided that the following conditions are met: 9 | 10 | * Redistributions of source code must retain the above copyright notice, this 11 | list of conditions and the following disclaimer. 12 | 13 | * Redistributions in binary form must reproduce the above copyright notice, 14 | this list of conditions and the following disclaimer in the documentation 15 | and/or other materials provided with the distribution. 16 | 17 | * Neither the name of the copyright holder nor the names of its 18 | contributors may be used to endorse or promote products derived from 19 | this software without specific prior written permission. 20 | 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 22 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 23 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 25 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 26 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 27 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 28 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 29 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 30 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TorchRec Ruby 2 | 3 | Deep learning recommendation systems for Ruby 4 | 5 | [![Build Status](https://github.com/ankane/torchrec-ruby/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/torchrec-ruby/actions) 6 | 7 | ## Installation 8 | 9 | Add this line to your application’s Gemfile: 10 | 11 | ```ruby 12 | gem "torchrec" 13 | ``` 14 | 15 | ## Getting Started 16 | 17 | This library follows the [Python API](https://pytorch.org/torchrec/). Many methods and options are missing at the moment. PRs welcome! 18 | 19 | ## Models 20 | 21 | DeepFM 22 | 23 | ```ruby 24 | TorchRec::Models::DeepFM::DenseArch.new(in_features, hidden_layer_size, embedding_dim) 25 | TorchRec::Models::DeepFM::OverArch.new(in_features) 26 | ``` 27 | 28 | DLRM 29 | 30 | ```ruby 31 | TorchRec::Models::DLRM::DenseArch.new(in_features, layer_sizes, device: nil) 32 | ``` 33 | 34 | ## Modules 35 | 36 | ```ruby 37 | TorchRec::Modules::Activation::SwishLayerNorm.new(input_dims, device: nil) 38 | TorchRec::Modules::CrossNet::CrossNet.new(in_features, num_layers) 39 | TorchRec::Modules::DeepFM::DeepFM.new(dense_module) 40 | TorchRec::Modules::DeepFM::FactorizationMachine.new 41 | TorchRec::Modules::MLP::MLP.new(in_size, layer_sizes, bias: true, activation: :relu, device: nil) 42 | TorchRec::Modules::MLP::Perceptron.new(in_size, out_size, bias: true, activation: Torch.method(:relu), device: nil) 43 | ``` 44 | 45 | ## History 46 | 47 | View the [changelog](https://github.com/ankane/torchrec-ruby/blob/master/CHANGELOG.md) 48 | 49 | ## Contributing 50 | 51 | Everyone is encouraged to help improve this project. Here are a few ways you can help: 52 | 53 | - [Report bugs](https://github.com/ankane/torchrec-ruby/issues) 54 | - Fix bugs and [submit pull requests](https://github.com/ankane/torchrec-ruby/pulls) 55 | - Write, clarify, or fix documentation 56 | - Suggest or add new features 57 | 58 | To get started with development: 59 | 60 | ```sh 61 | git clone https://github.com/ankane/torchrec-ruby.git 62 | cd torchrec-ruby 63 | bundle install 64 | bundle exec rake test 65 | ``` 66 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require "bundler/gem_tasks" 2 | require "rake/testtask" 3 | 4 | Rake::TestTask.new(:test) do |t| 5 | t.libs << "test" 6 | t.pattern = "test/**/*_test.rb" 7 | end 8 | 9 | task default: :test 10 | -------------------------------------------------------------------------------- /lib/torchrec.rb: -------------------------------------------------------------------------------- 1 | # dependencies 2 | require "torch-rb" 3 | 4 | # models 5 | require_relative "torchrec/models/deepfm/dense_arch" 6 | require_relative "torchrec/models/deepfm/over_arch" 7 | require_relative "torchrec/models/dlrm/dense_arch" 8 | 9 | # modules 10 | require_relative "torchrec/modules/activation/swish_layer_norm" 11 | require_relative "torchrec/modules/cross_net/cross_net" 12 | require_relative "torchrec/modules/deepfm/deepfm" 13 | require_relative "torchrec/modules/deepfm/factorization_machine" 14 | require_relative "torchrec/modules/mlp/mlp" 15 | require_relative "torchrec/modules/mlp/perceptron" 16 | require_relative "torchrec/modules/utils" 17 | 18 | # sparse 19 | require_relative "torchrec/sparse/jagged_tensor" 20 | 21 | # other 22 | require_relative "torchrec/version" 23 | 24 | module TorchRec 25 | class Error < StandardError; end 26 | end 27 | -------------------------------------------------------------------------------- /lib/torchrec/models/deepfm/dense_arch.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Models 3 | module DeepFM 4 | class DenseArch < Torch::NN::Module 5 | def initialize(in_features, hidden_layer_size, embedding_dim) 6 | super() 7 | @model = Torch::NN::Sequential.new( 8 | Torch::NN::Linear.new(in_features, hidden_layer_size), 9 | Torch::NN::ReLU.new, 10 | Torch::NN::Linear.new(hidden_layer_size, embedding_dim), 11 | Torch::NN::ReLU.new 12 | ) 13 | end 14 | 15 | def forward(features) 16 | @model.call(features) 17 | end 18 | end 19 | end 20 | end 21 | end 22 | -------------------------------------------------------------------------------- /lib/torchrec/models/deepfm/over_arch.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Models 3 | module DeepFM 4 | class OverArch < Torch::NN::Module 5 | def initialize(in_features) 6 | super() 7 | @model = Torch::NN::Sequential.new( 8 | Torch::NN::Linear.new(in_features, 1), 9 | Torch::NN::Sigmoid.new 10 | ) 11 | end 12 | 13 | def forward(features) 14 | @model.call(features) 15 | end 16 | end 17 | end 18 | end 19 | end 20 | -------------------------------------------------------------------------------- /lib/torchrec/models/dlrm/dense_arch.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Models 3 | module DLRM 4 | class DenseArch < Torch::NN::Module 5 | def initialize(in_features, layer_sizes, device: nil) 6 | super() 7 | @model = Modules::MLP::MLP.new( 8 | in_features, layer_sizes, bias: true, activation: :relu, device: device 9 | ) 10 | end 11 | 12 | def forward(features) 13 | @model.call(features) 14 | end 15 | end 16 | end 17 | end 18 | end 19 | -------------------------------------------------------------------------------- /lib/torchrec/modules/activation/swish_layer_norm.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module Activation 4 | class SwishLayerNorm < Torch::NN::Module 5 | def initialize(input_dims, device: nil) 6 | super() 7 | @norm = Torch::NN::Sequential.new( 8 | # TODO add device 9 | Torch::NN::LayerNorm.new(input_dims), #, device: device), 10 | Torch::NN::Sigmoid.new 11 | ) 12 | end 13 | 14 | def forward(input) 15 | input * @norm.call(input) 16 | end 17 | end 18 | end 19 | end 20 | end 21 | -------------------------------------------------------------------------------- /lib/torchrec/modules/cross_net/cross_net.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module CrossNet 4 | class CrossNet < Torch::NN::Module 5 | def initialize(in_features, num_layers) 6 | super() 7 | @num_layers = num_layers 8 | @kernels = Torch::NN::ParameterList.new( 9 | @num_layers.times.map do |i| 10 | Torch::NN::Parameter.new( 11 | Torch::NN::Init.xavier_normal!(Torch.empty(in_features, in_features)) 12 | ) 13 | end 14 | ) 15 | @bias = Torch::NN::ParameterList.new( 16 | @num_layers.times.map do |i| 17 | Torch::NN::Parameter.new(Torch::NN::Init.zeros!(Torch.empty(in_features, 1))) 18 | end 19 | ) 20 | end 21 | 22 | def forward(input) 23 | x_0 = input.unsqueeze(2) # (B, N, 1) 24 | x_l = x_0 25 | 26 | @num_layers.times do |layer| 27 | xl_w = Torch.matmul(@kernels[layer], x_l) # (B, N, 1) 28 | x_l = x_0 * (xl_w + @bias[layer]) + x_l # (B, N, 1) 29 | end 30 | 31 | Torch.squeeze(x_l, dim: 2) 32 | end 33 | end 34 | end 35 | end 36 | end 37 | -------------------------------------------------------------------------------- /lib/torchrec/modules/deepfm/deepfm.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module DeepFM 4 | class DeepFM < Torch::NN::Module 5 | def initialize(dense_module) 6 | super() 7 | @dense_module = dense_module 8 | end 9 | 10 | def forward(embeddings) 11 | deepfm_input = flatten_input(embeddings) 12 | @dense_module.call(deepfm_input) 13 | end 14 | 15 | private 16 | 17 | def flatten_input(inputs) 18 | Torch.cat(inputs.map { |input| input.flatten(1) }, dim: 1) 19 | end 20 | end 21 | end 22 | end 23 | end 24 | -------------------------------------------------------------------------------- /lib/torchrec/modules/deepfm/factorization_machine.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module DeepFM 4 | class FactorizationMachine < Torch::NN::Module 5 | def initialize 6 | super() 7 | end 8 | 9 | def forward(embeddings) 10 | fm_input = flatten_input(embeddings) 11 | sum_of_input = Torch.sum(fm_input, dim: 1, keepdim: true) 12 | sum_of_square = Torch.sum(fm_input * fm_input, dim: 1, keepdim: true) 13 | square_of_sum = sum_of_input * sum_of_input 14 | cross_term = square_of_sum - sum_of_square 15 | cross_term = Torch.sum(cross_term, dim: 1, keepdim: true) * 0.5 # [B, 1] 16 | cross_term 17 | end 18 | 19 | private 20 | 21 | def flatten_input(inputs) 22 | Torch.cat(inputs.map { |input| input.flatten(1) }, dim: 1) 23 | end 24 | end 25 | end 26 | end 27 | end 28 | -------------------------------------------------------------------------------- /lib/torchrec/modules/mlp/mlp.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module MLP 4 | class MLP < Torch::NN::Module 5 | def initialize(in_size, layer_sizes, bias: true, activation: :relu, device: nil) 6 | super() 7 | 8 | if activation == :relu 9 | activation = Torch.method(:relu) 10 | elsif activation == :sigmoid 11 | activation = Torch.method(:sigmoid) 12 | end 13 | 14 | if !activation.is_a?(Symbol) 15 | @mlp = Torch::NN::Sequential.new( 16 | *layer_sizes.length.times.map do |i| 17 | Perceptron.new( 18 | i > 0 ? layer_sizes[i - 1] : in_size, 19 | layer_sizes[i], 20 | bias: bias, 21 | activation: Utils.extract_module_or_tensor_callable(activation), 22 | device: device 23 | ) 24 | end 25 | ) 26 | else 27 | if activation == :swish_layernorm 28 | @mlp = Torch::NN::Sequential.new( 29 | *layer_sizes.length.times.map do |i| 30 | Perceptron.new( 31 | i > 0 ? layer_sizes[i - 1] : in_size, 32 | layer_sizes[i], 33 | bias: bias, 34 | activation: Activation::SwishLayerNorm.new(layer_sizes[i], device: device), 35 | device: device 36 | ) 37 | end 38 | ) 39 | else 40 | raise ArgumentError, "This MLP only supports activation function of :relu, :sigmoid, and :swish_layernorm" 41 | end 42 | end 43 | end 44 | 45 | def forward(input) 46 | @mlp.call(input) 47 | end 48 | end 49 | end 50 | end 51 | end 52 | -------------------------------------------------------------------------------- /lib/torchrec/modules/mlp/perceptron.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module MLP 4 | class Perceptron < Torch::NN::Module 5 | def initialize(in_size, out_size, bias: true, activation: Torch.method(:relu), device: nil) 6 | super() 7 | @out_size = out_size 8 | @in_size = in_size 9 | @linear = Torch::NN::Linear.new( 10 | # TODO add device 11 | @in_size, @out_size, bias: bias #, device: device 12 | ) 13 | @activation_fn = activation 14 | end 15 | 16 | def forward(input) 17 | @activation_fn.call(@linear.call(input)) 18 | end 19 | end 20 | end 21 | end 22 | end 23 | -------------------------------------------------------------------------------- /lib/torchrec/modules/utils.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Modules 3 | module Utils 4 | class << self 5 | # TODO 6 | def extract_module_or_tensor_callable(module_or_callable) 7 | module_or_callable 8 | end 9 | end 10 | end 11 | end 12 | end 13 | -------------------------------------------------------------------------------- /lib/torchrec/sparse/jagged_tensor.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | module Sparse 3 | class JaggedTensor 4 | def initialize(values, weights: nil, lengths: nil, offsets: nil) 5 | @values = values 6 | @weights = weights 7 | assert_offsets_or_lengths_is_provided(offsets, lengths) 8 | if !offsets.nil? 9 | assert_tensor_has_no_elements_or_has_integers(offsets, "offsets") 10 | end 11 | if !lengths.nil? 12 | assert_tensor_has_no_elements_or_has_integers(lengths, "lengths") 13 | end 14 | @lengths = lengths 15 | @offsets = offsets 16 | end 17 | 18 | private 19 | 20 | def assert_offsets_or_lengths_is_provided(offsets, lengths) 21 | if offsets.nil? && lengths.nil? 22 | raise ArgumentError, "Must provide lengths or offsets" 23 | end 24 | end 25 | 26 | def assert_tensor_has_no_elements_or_has_integers(tensor, tensor_name) 27 | if tensor.numel != 0 && ![:int64, :int32, :int16, :int8, :uint8].include?(tensor.dtype) 28 | raise ArgumentError, "#{tensor_name} must be of integer type, but got #{tensor.dtype}" 29 | end 30 | end 31 | end 32 | end 33 | end 34 | -------------------------------------------------------------------------------- /lib/torchrec/version.rb: -------------------------------------------------------------------------------- 1 | module TorchRec 2 | VERSION = "0.0.4" 3 | end 4 | -------------------------------------------------------------------------------- /test/models_test.rb: -------------------------------------------------------------------------------- 1 | require_relative "test_helper" 2 | 3 | class ModelsTest < Minitest::Test 4 | def test_deepfm_dense_arch 5 | dense_arch = TorchRec::Models::DeepFM::DenseArch.new(10, 10, 3) 6 | dense_embedded = dense_arch.call(Torch.rand([20, 10])) 7 | assert_equal [20, 3], dense_embedded.shape 8 | end 9 | 10 | def test_deepfm_over_arch 11 | over_arch = TorchRec::Models::DeepFM::OverArch.new(10) 12 | logits = over_arch.call(Torch.rand([20, 10])) 13 | assert_equal [20, 1], logits.shape 14 | end 15 | 16 | def test_dlrm_dense_arch 17 | dense_arch = TorchRec::Models::DLRM::DenseArch.new(10, [15, 3]) 18 | dense_embedded = dense_arch.call(Torch.rand([20, 10])) 19 | assert_equal [20, 3], dense_embedded.shape 20 | end 21 | end 22 | -------------------------------------------------------------------------------- /test/modules_test.rb: -------------------------------------------------------------------------------- 1 | require_relative "test_helper" 2 | 3 | class ModulesTest < Minitest::Test 4 | def test_swish_layer_norm 5 | sln = TorchRec::Modules::Activation::SwishLayerNorm.new(100) 6 | assert sln 7 | end 8 | 9 | def test_cross_net 10 | batch_size = 3 11 | num_layers = 2 12 | in_features = 10 13 | input = Torch.randn(batch_size, in_features) 14 | dcn = TorchRec::Modules::CrossNet::CrossNet.new(in_features, num_layers) 15 | output = dcn.call(input) 16 | assert output 17 | end 18 | 19 | def test_deepfm 20 | batch_size = 3 21 | output_dim = 30 22 | input_embeddings = [ 23 | Torch.randn(batch_size, 2, 64), 24 | Torch.randn(batch_size, 2, 32) 25 | ] 26 | dense_module = Torch::NN::Linear.new(192, output_dim) 27 | deepfm = TorchRec::Modules::DeepFM::DeepFM.new(dense_module) 28 | deep_fm_output = deepfm.call(input_embeddings) 29 | assert deep_fm_output 30 | end 31 | 32 | def test_factorization_machine 33 | batch_size = 3 34 | input_embeddings = [ 35 | Torch.randn(batch_size, 2, 64), 36 | Torch.randn(batch_size, 2, 32) 37 | ] 38 | fm = TorchRec::Modules::DeepFM::FactorizationMachine.new 39 | output = fm.call(input_embeddings) 40 | assert output 41 | end 42 | 43 | def test_mlp 44 | batch_size = 3 45 | in_size = 40 46 | input = Torch.randn(batch_size, in_size) 47 | 48 | layer_sizes = [16, 8, 4] 49 | mlp_module = TorchRec::Modules::MLP::MLP.new(in_size, layer_sizes, bias: true) 50 | output = mlp_module.call(input) 51 | assert_equal output.shape, [batch_size, layer_sizes[-1]] 52 | end 53 | 54 | def test_perceptron 55 | batch_size = 3 56 | in_size = 40 57 | input = Torch.randn(batch_size, in_size) 58 | 59 | out_size = 16 60 | perceptron = TorchRec::Modules::MLP::Perceptron.new(in_size, out_size, bias: true) 61 | 62 | output = perceptron.call(input) 63 | assert_equal output.shape, [batch_size, out_size] 64 | end 65 | end 66 | -------------------------------------------------------------------------------- /test/sparse_test.rb: -------------------------------------------------------------------------------- 1 | require_relative "test_helper" 2 | 3 | class SparseTest < Minitest::Test 4 | def test_jagged_tensor 5 | assert TorchRec::Sparse::JaggedTensor.new(Torch.tensor([]), offsets: Torch.tensor([])) 6 | end 7 | end 8 | -------------------------------------------------------------------------------- /test/test_helper.rb: -------------------------------------------------------------------------------- 1 | require "bundler/setup" 2 | Bundler.require(:default) 3 | require "minitest/autorun" 4 | require "minitest/pride" 5 | -------------------------------------------------------------------------------- /torchrec.gemspec: -------------------------------------------------------------------------------- 1 | require_relative "lib/torchrec/version" 2 | 3 | Gem::Specification.new do |spec| 4 | spec.name = "torchrec" 5 | spec.version = TorchRec::VERSION 6 | spec.summary = "Deep learning recommendation systems for Ruby" 7 | spec.homepage = "https://github.com/ankane/torchrec-ruby" 8 | spec.license = "BSD-3-Clause" 9 | 10 | spec.author = "Andrew Kane" 11 | spec.email = "andrew@ankane.org" 12 | 13 | spec.files = Dir["*.{md,txt}", "{lib}/**/*"] 14 | spec.require_path = "lib" 15 | 16 | spec.required_ruby_version = ">= 3.1" 17 | 18 | spec.add_dependency "torch-rb", ">= 0.13" 19 | end 20 | --------------------------------------------------------------------------------