transformers-rb 0.1.4 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 16b409e954f6bcc45fd4f3f2db94dc92c87d47c4c06936162978cc7d7e54fc09
-  data.tar.gz: c40ea58b7531e89a041ce2782dea89536f85227903e2dc9a60113afe041bb9f7
+  metadata.gz: 815a531b2876ca2f66d9df65e13a08f42ab0b4a98d0d1ef459961fc3069faaa7
+  data.tar.gz: 842158dd34cd785efff7a6df7524acd019afc6ead7738cfd9ddf6efadb89d97d
 SHA512:
-  metadata.gz: 0576500ca9fe9379aae4c2cc050aa34c90eea7b2d5251b6139c48d88b5107086e6197aad5210bc74d7969a5bfa0458d31ce7faf561df69b2ae9b2a6400280ce0
-  data.tar.gz: b5b0a865acdd37bcde11571e365a1d39b473f3dcbbe737a082f37002f886fe38f0f57c77f44236f859939d1d3b1df9659bb524a483cd0b9048dbc4cd472a3355
+  metadata.gz: 2a4fa39c9d32b4ba83eeafe128d10af0873a910807ad67870359e8bffcdf94bdeeb18dea8bc0b2372926ded909e8223a2603388297aad357bca26ea0c6eb0baa
+  data.tar.gz: bfac7d5c91e91bbbf1770ea71a9a9c69bca71c1e0ded3b0ca696156c99c331a5988879a3cf48731145db7be1d1c1b14dc7e573b728e1d1a123faa251ab56a560

data/CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.1.5 (2024-11-01)
+
+- Fixed error with pipelines when called more than 10 times
+- Fixed `device` option for pipelines
+- Fixed error with `reranking` pipeline
+
 ## 0.1.4 (2024-10-22)
 
 - Added `BertForSequenceClassification`

data/lib/transformers/modeling_utils.rb

@@ -56,6 +56,10 @@ module Transformers
   end
 
   module ModuleUtilsMixin
+    def device
+      get_parameter_device(self)
+    end
+
     def get_extended_attention_mask(
       attention_mask,
       input_shape,
@@ -112,6 +116,13 @@ module Transformers
 
       head_mask
     end
+
+    private
+
+    def get_parameter_device(parameter)
+      # TODO return Torch::Device in Torch.rb
+      Torch.device(parameter.parameters[0].device)
+    end
   end
 
   class PreTrainedModel < Torch::NN::Module

data/lib/transformers/models/xlm_roberta/modeling_xlm_roberta.rb

@@ -1213,4 +1213,6 @@ module Transformers
       end
     end
   end
+
+  XLMRobertaForSequenceClassification = XlmRoberta::XLMRobertaForSequenceClassification
 end

data/lib/transformers/pipelines/base.rb

@@ -123,6 +123,22 @@ module Transformers
         end
       end
 
+      if @framework == "pt"
+        if device.is_a?(String)
+          @device = Torch.device(device)
+        else
+          # TODO update default in 0.2.0
+          @device = @model.device
+        end
+      else
+        raise Todo
+      end
+
+      # TODO Fix eql? for Torch::Device in Torch.rb
+      if @device.type != @model.device.type || @device.index != @model.device.index
+        @model.to(@device)
+      end
+
       @call_count = 0
       @batch_size = kwargs.delete(:batch_size)
       @num_workers = kwargs.delete(:num_workers)
@@ -133,6 +149,24 @@ module Transformers
       @model.dtype
     end
 
+    def _ensure_tensor_on_device(inputs, device)
+      # TODO move
+      inputs = inputs.to_h if inputs.is_a?(BatchEncoding)
+
+      if inputs.is_a?(ModelOutput)
+        inputs.instance_variable_get(:@data).transform_values! { |tensor| _ensure_tensor_on_device(tensor, device) }
+        inputs
+      elsif inputs.is_a?(Hash)
+        inputs.transform_values { |tensor| _ensure_tensor_on_device(tensor, device) }
+      elsif inputs.is_a?(Array)
+        inputs.map { |item| _ensure_tensor_on_device(item, device) }
+      elsif inputs.is_a?(Torch::Tensor)
+        inputs.to(device)
+      else
+        inputs
+      end
+    end
+
     def check_model_type(supported_models)
       if !supported_models.is_a?(Array)
         supported_models_names = []
@@ -245,7 +279,10 @@ module Transformers
     end
 
     def forward(model_inputs, **forward_params)
-      _forward(model_inputs, **forward_params)
+      model_inputs = _ensure_tensor_on_device(model_inputs, @device)
+      model_outputs = _forward(model_inputs, **forward_params)
+      model_outputs = _ensure_tensor_on_device(model_outputs, Torch.device("cpu"))
+      model_outputs
     end
 
     def run_single(inputs, preprocess_params, forward_params, postprocess_params)

data/lib/transformers/pipelines/reranking.rb

@@ -8,7 +8,8 @@ module Transformers
       @tokenizer.(
         [inputs[:query]] * inputs[:documents].length,
         text_pair: inputs[:documents],
-        return_tensors: @framework
+        return_tensors: @framework,
+        padding: true
       )
     end
 

data/lib/transformers/version.rb

@@ -1,3 +1,3 @@
 module Transformers
-  VERSION = "0.1.4"
+  VERSION = "0.1.5"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: transformers-rb
 version: !ruby/object:Gem::Version
-  version: 0.1.4
+  version: 0.1.5
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-10-22 00:00:00.000000000 Z
+date: 2024-11-01 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: numo-narray