rbbt-dm 1.2.9 → 1.2.10

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: db2cbab94e21fd2ca67f7306fa9941b59cbfb2865382e5439edf6313f50309e7
-  data.tar.gz: f4acf3651daa90ef23bc454c62df68e208976a977d51e2e85d02558d48897187
+  metadata.gz: 9d53609453e1c3bd589c95071569583bff5f11224a200850c0d9e85775e5a2ce
+  data.tar.gz: 501f436caff07c990c09adec4caa60d32fe577b35209a65e8f96418ab2acc422
 SHA512:
-  metadata.gz: 7786759636450821aabca306cd210ab3d201e094b81bb70052d57d7bfb6e4de73a198576fe4b002487baf7997138f9c53b91644632cd12cc79b40ff62141a70a
-  data.tar.gz: 9870745068a897909170f3a6187e520e6530b121f8ad4ab40224c3369a16a8bb1c1e55bb8ca8fd892943ec1fcead9f4661e49c3d97d07d3740236a1ec4f69a34
+  metadata.gz: 4a53ccc8a5eac633e344beffd36073e1828d848d9743552f12e75a54cbda337050152b72fd95d7d97d4aeed50f811dcc8475c97bc292ba9bcdd57e9d32f0e91b
+  data.tar.gz: 30fef4ee2c023e141c0ef991ae35d1a0913d4444138fda778760b821ccc48d85d67dc44d4c251d7f84e3b1bbe6ecb6f3b64aee13fbdc443cb33bbed2b34a586a

data/lib/rbbt/vector/model/huggingface.rb

@@ -2,53 +2,21 @@ require 'rbbt/vector/model/torch'
 
 class HuggingfaceModel < TorchModel
 
-  def self.tsv_dataset(tsv_dataset_file, elements, labels = nil, class_labels = nil)
-
-    if labels
-      labels = case class_labels
-               when Array
-                 labels.collect{|l| class_labels.index l}
-               when Hash
-                 inverse_class_labels = {}
-                 class_labels.each{|c,l| inverse_class_labels[l] = c }
-                 labels.collect{|l| inverse_class_labels[l]}
-               else
-                 labels
-               end
-
-      Open.write(tsv_dataset_file) do |ffile|
-        ffile.puts ["label", "text"].flatten * "\t"
-        elements.zip(labels).each do |element,label|
-          element = element.gsub("\n", " ")
-          ffile.puts [label, element].flatten * "\t"
-        end
-        ffile.sync
-      end
-    else
-      Open.write(tsv_dataset_file) do |ffile|
-        ffile.puts ["text"].flatten * "\t"
-        elements.each do |element|
-          element = element.gsub("\n", " ")
-          ffile.puts element 
-        end
-        ffile.sync
-      end
-    end
-
-    tsv_dataset_file
-  end
-
   def initialize(task, checkpoint, dir = nil, model_options = {})
-    super(dir, model_options)
+    super(dir, nil, model_options)
 
     @model_options = Misc.add_defaults @model_options, :task => task, :checkpoint => checkpoint
 
     init_model do 
       checkpoint = @model_path && File.directory?(@model_path) ? @model_path : @model_options[:checkpoint]
+
       model = RbbtPython.call_method("rbbt_dm.huggingface", :load_model, 
                                      @model_options[:task], checkpoint, **(IndiferentHash.setup(model_options[:model_args]) || {}))
+
+      tokenizer_checkpoint = @model_options[:tokenizer_checkpoint] || checkpoint
+
       tokenizer = RbbtPython.call_method("rbbt_dm.huggingface", :load_tokenizer, 
-                                         @model_options[:task], checkpoint, **(IndiferentHash.setup(model_options[:tokenizer_args]) || {}))
+                                         @model_options[:task], tokenizer_checkpoint, **(IndiferentHash.setup(model_options[:tokenizer_args]) || {}))
 
       [model, tokenizer]
     end
@@ -75,7 +43,7 @@ class HuggingfaceModel < TorchModel
           checkpoint_dir = File.join(tmpdir, 'checkpoints')
         end
 
-        dataset_file = HuggingfaceModel.tsv_dataset(tsv_file, texts)
+        dataset_file = TorchModel.text_dataset(tsv_file, texts)
         training_args_obj = RbbtPython.call_method("rbbt_dm.huggingface", :training_args, checkpoint_dir, @model_options[:training_args])
 
         begin
@@ -102,7 +70,7 @@ class HuggingfaceModel < TorchModel
       end
 
       training_args_obj = RbbtPython.call_method("rbbt_dm.huggingface", :training_args, checkpoint_dir, @model_options[:training_args])
-      dataset_file = HuggingfaceModel.tsv_dataset(tsv_file, texts, labels, @model_options[:class_labels])
+      dataset_file = HuggingfaceModel.text_dataset(tsv_file, texts, labels, @model_options[:class_labels])
 
       RbbtPython.call_method("rbbt_dm.huggingface", :train_model, model, tokenizer, training_args_obj, dataset_file, @model_options[:class_weights])
 

data/lib/rbbt/vector/model/python.rb

@@ -0,0 +1,33 @@
+require 'rbbt/vector/model'
+require 'rbbt/util/python'
+
+RbbtPython.add_path Rbbt.python.find(:lib)
+RbbtPython.init_rbbt
+
+class PythonModel < VectorModel
+  attr_accessor :python_class, :python_module
+  def initialize(dir, python_class = nil, python_module = nil, model_options = nil)
+    python_module = :model if python_module.nil?
+    model_options, python_module = python_module, :model if model_options.nil? && Hash === python_module
+    model_options = {} if model_options.nil?
+
+    super(dir, model_options)
+
+    @python_class = python_class
+    @python_module = python_module
+
+    init_model do
+      RbbtPython.add_path @directory 
+      RbbtPython.class_new_obj(@python_module, @python_class, **model_options)
+    end if python_class
+
+    eval_model do |features,list=false|
+      init
+      if list
+        model.eval(features)
+      else
+        model.eval([features])[0]
+      end
+    end
+  end
+end

data/lib/rbbt/vector/model/pytorch_lightning.rb

@@ -2,34 +2,30 @@ require 'rbbt/vector/model/torch'
 
 class PytorchLightningModel < TorchModel
   attr_accessor :loader, :val_loader, :trainer
-  def initialize(module_name, class_name, dir = nil, model_options = {})
-    super(dir, model_options)
-    @module_name = module_name
-    @class_name = class_name
-
-    init_model do 
-      RbbtPython.pyimport @module_name
-      RbbtPython.class_new_obj(@module_name, @class_name, @model_options[:model_args] || {})
-    end
+  def initialize(...)
+    super(...)
 
     train_model do |features,labels|
       model = init
-      raise "Use the loader" if @loader.nil?
-      raise "Use the trainer" if @trainer.nil?
-      
-      trainer.fit(model, @loader, @val_loader)
-    end
-
-    eval_model do |features,list|
-      if list
-        model.call(RbbtPython.call_method(:torch, :tensor, features))
-      else
-        model.call(RbbtPython.call_method(:torch, :tensor, [features]))
+      loader = self.loader
+      val_loader = self.val_loader
+      if (features && features.any?) && loader.nil?
+        TmpFile.with_file do |tsv_dataset_file|
+          TorchModel.feature_dataset(tsv_dataset_file, features, labels)
+          RbbtPython.pyimport :rbbt_dm
+          loader = RbbtPython.rbbt_dm.tsv(tsv_dataset_file)
+        end
       end
+      trainer.fit(model, loader, val_loader)
+      TorchModel.save_architecture(model, model_path) if @directory
+      TorchModel.save_state(model, model_path) if @directory
     end
-
   end
-end
 
-if __FILE__ == $0
+  def trainer
+    @trainer ||= begin
+                   options = @model_options[:training_args] || @model_options[:trainer_args] 
+                   RbbtPython.class_new_obj("pytorch_lightning", "Trainer", options || {})
+                 end
+  end
 end

data/lib/rbbt/vector/model/torch/dataloader.rb

@@ -0,0 +1,58 @@
+class TorchModel
+  def self.feature_tsv(elements, labels = nil, class_labels = nil)
+    tsv = TSV.setup({}, :key_field => "ID", :fields => ["features"], :type => :flat)
+    if labels
+      tsv.fields = tsv.fields + ["label"]
+      labels = case class_labels
+               when Array
+                 labels.collect{|l| class_labels.index l}
+               when Hash
+                 inverse_class_labels = {}
+                 class_labels.each{|c,l| inverse_class_labels[l] = c }
+                 labels.collect{|l| inverse_class_labels[l]}
+               else
+                 labels
+               end
+      elements.zip(labels).each_with_index do |p,i|
+        features, label = p
+        id = i
+        if Array === features
+          tsv[id] = features + [label]
+        else
+          tsv[id] = [features, label]
+        end
+      end
+    else
+      elements.each_with_index do |features,i|
+        id = i
+        if Array === features
+          tsv[id] = features
+        else
+          tsv[id] = [features]
+        end
+      end
+    end
+    tsv
+  end
+
+  def self.feature_dataset(tsv_dataset_file, elements, labels = nil, class_labels = nil)
+    tsv = feature_tsv(elements, labels, class_labels)
+    Open.write(tsv_dataset_file, tsv.to_s)
+    tsv_dataset_file
+  end
+
+  def self.text_dataset(tsv_dataset_file, elements, labels = nil, class_labels = nil)
+    elements = elements.collect{|e| e.gsub("\n", ' ') }
+    tsv = feature_tsv(elements, labels, class_labels)
+    if labels.nil?
+      tsv.fields[0] = "text"
+      tsv.type = :single
+    else
+      tsv.fields[0] = "text"
+      tsv.type = :list
+    end
+    Open.write(tsv_dataset_file, tsv.to_s)
+    tsv_dataset_file
+  end
+
+end

data/lib/rbbt/vector/model/torch/helpers.rb

@@ -0,0 +1,52 @@
+class TorchModel
+  module Tensor
+    def to_ruby
+      RbbtPython.numpy2ruby(self)
+    end
+    def self.setup(obj)
+      obj.extend Tensor
+    end
+  end
+
+  def self.init_python
+    RbbtPython.pyimport :torch
+    RbbtPython.pyimport :rbbt
+    RbbtPython.pyimport :rbbt_dm
+    RbbtPython.pyfrom :rbbt_dm, import: :util
+    RbbtPython.pyfrom :torch, import: :nn
+  end
+
+  def self.optimizer(model, training_args)
+    begin
+      learning_rate = training_args[:learning_rate] || 0.01
+      RbbtPython.torch.optim.SGD.new(model.parameters(), lr: learning_rate)
+    end
+  end
+
+  def self.device(model_options)
+    case model_options[:device]
+    when String, Symbol
+      RbbtPython.torch.device(model_options[:device].to_s)
+    when nil
+      RbbtPython.rbbt_dm.util.device()
+    else
+        model_options[:device]
+    end
+  end
+
+  def self.dtype(model_options)
+    case model_options[:dtype]
+    when String, Symbol
+      RbbtPython.torch.call(model_options[:dtype])
+    when nil
+      nil
+    else
+      model_options[:dtype]
+    end
+  end
+
+  def self.tensor(obj, device, dtype)
+    RbbtPython.torch.tensor(obj, dtype: dtype, device: device)
+  end
+
+end

data/lib/rbbt/vector/model/torch/introspection.rb

@@ -0,0 +1,31 @@
+class TorchModel
+  def self.get_layer(model, layer = nil)
+    if layer.nil?
+      model
+    else
+      layer.split(".").inject(model){|acc,l| PyCall.getattr(acc, l.to_sym) }
+    end
+  end
+  def get_layer(...); TorchModel.get_layer(model, ...); end
+
+  def self.get_weights(model, layer = nil)
+    Tensor.setup PyCall.getattr(get_layer(model, layer), :weight)
+  end
+  def get_weights(...); TorchModel.get_weights(model, ...); end
+
+  def self.freeze(layer)
+    begin
+      PyCall.getattr(layer, :weight).requires_grad = false
+    rescue
+    end
+    RbbtPython.iterate(layer.children) do |layer|
+      freeze(layer)
+    end
+  end
+  def self.freeze_layer(model, layer)
+    layer = get_layer(model, layer)
+    freeze(layer)
+  end
+  def freeze_layer(...); TorchModel.freeze_layer(model, ...); end
+
+end

data/lib/rbbt/vector/model/torch/load_and_save.rb

@@ -0,0 +1,30 @@
+class TorchModel
+  def self.model_architecture(model_path)
+    model_path + '.architecture'
+  end
+
+  def self.save_state(model, model_path)
+    Log.debug "Saving model state into #{model_path}"
+    RbbtPython.torch.save(model.state_dict(), model_path)
+  end
+
+  def self.load_state(model, model_path)
+    return model unless Open.exists?(model_path)
+    Log.debug "Loading model state from #{model_path}"
+    model.load_state_dict(RbbtPython.torch.load(model_path))
+    model
+  end
+
+  def self.save_architecture(model, model_path)
+    model_architecture = model_architecture(model_path)
+    Log.debug "Saving model architecture into #{model_architecture}"
+    RbbtPython.torch.save(model, model_architecture)
+  end
+
+  def self.load_architecture(model_path)
+    model_architecture = model_architecture(model_path)
+    return unless Open.exists?(model_architecture)
+    Log.debug "Loading model architecture from #{model_architecture}"
+    RbbtPython.torch.load(model_architecture)
+  end
+end

data/lib/rbbt/vector/model/torch.rb

@@ -1,37 +1,71 @@
-require 'rbbt/vector/model'
-require 'rbbt/util/python'
+require_relative 'python'
 
-RbbtPython.add_path Rbbt.python.find(:lib)
-RbbtPython.init_rbbt
+class TorchModel < PythonModel
 
-class TorchModel < VectorModel
+  attr_accessor :model, :criterion, :optimizer, :training_args
 
-  attr_accessor :model
+  def initialize(...)
+    TorchModel.init_python
+    super(...)
+    @training_args = model_options[:training_args] || {}
 
-  def self.get_layer(model, layer)
-    layer.split(".").inject(model){|acc,l| PyCall.getattr(acc, l.to_sym) }
-  end
+    init_model do
+      model = TorchModel.load_architecture(model_path) 
+      if model.nil?
+        RbbtPython.add_path @directory 
+        RbbtPython.class_new_obj(@python_module, @python_class, **model_options)
+      else
+        TorchModel.load_state(model, model_path)
+      end
+    end
 
-  def self.get_weights(model, layer)
-    PyCall.getattr(get_layer(model, layer), :weight)
-  end
+    eval_model do |features,list=false|
+      init
+      @device ||= TorchModel.device(model_options)
+      @dtype ||= TorchModel.dtype(model_options)
+      model.to(@device)
 
-  def self.freeze(layer)
-    begin
-      PyCall.getattr(layer, :weight).requires_grad = false
-    rescue
-    end
-    RbbtPython.iterate(layer.children) do |layer|
-      freeze(layer)
+      tensor = list ? TorchModel.tensor(features, @device, @dtype) : TorchModel.tensor([features], @device, @dtype)
+
+      loss, res = model.call(tensor)
+
+      res = loss if res.nil?
+
+      res = TorchModel::Tensor.setup(list ? res : res[0])
+
+      res
     end
-  end
 
-  def self.freeze_layer(model, layer)
-    layer = get_layer(model, layer)
-    freeze(layer)
-  end
+    train_model do |features,labels|
+      init
+      @device ||= TorchModel.device(model_options)
+      @dtype ||= TorchModel.dtype(model_options)
+      model.to(@device)
+      @optimizer ||= TorchModel.optimizer(model, training_args)
+      epochs = training_args[:epochs] || 3
+
+      inputs = TorchModel.tensor(features, @device, @dtype)
+      #target = TorchModel.tensor(labels.collect{|v| [v] }, @device, @dtype)
+      target = TorchModel.tensor(labels, @device, @dtype)
 
-  def initialize(dir, model_options = {})
-    super(dir, model_options)
+      Log::ProgressBar.with_bar epochs, :desc => "Training" do |bar|
+        epochs.times do |i|
+          @optimizer.zero_grad()
+          outputs = model.call(inputs)
+          outputs = outputs.squeeze() if target.dim() == 1
+          loss = criterion.call(outputs, target)
+          loss.backward()
+          @optimizer.step
+          Log.debug "Epoch #{i}, loss #{loss}"
+          bar.tick
+        end
+      end
+      TorchModel.save_architecture(model, model_path) if @directory
+      TorchModel.save_state(model, model_path) if @directory
+    end
   end
 end
+require_relative 'torch/helpers'
+require_relative 'torch/dataloader'
+require_relative 'torch/introspection'
+require_relative 'torch/load_and_save'

data/lib/rbbt/vector/model.rb

@@ -448,10 +448,10 @@ cat(paste(label, sep="\\n", collapse="\\n"));
         end
 
         test_set = feature_folds[fix]
-        train_set = feature_folds.values_at(*rest).inject([]){|acc,e| acc += e; acc}
+        train_set = feature_folds.values_at(*rest).flatten(1)
 
         test_labels = labels_folds[fix]
-        train_labels = labels_folds.values_at(*rest).flatten
+        train_labels = labels_folds.values_at(*rest).flatten(1)
 
         @features = train_set
         @labels = train_labels

data/python/rbbt_dm/__init__.py

@@ -1,6 +1,8 @@
-from torch.utils.data import Dataset, DataLoader
+import rbbt
+import torch
+from .util import *
 
-class TSVDataset(Dataset):
+class TSVDataset(torch.utils.data.Dataset):
     def __init__(self, tsv):
         self.tsv = tsv
 
@@ -20,29 +22,10 @@ class TSVDataset(Dataset):
         return len(self.tsv)
 
 def tsv_dataset(filename, *args, **kwargs):
-    import rbbt
     return TSVDataset(rbbt.tsv(filename, *args, **kwargs))
 
 def tsv(*args, **kwargs):
     return tsv_dataset(*args, **kwargs)
 
 def data_dir():
-    import rbbt
     return rbbt.path('var/rbbt_dm/data')
-
-if __name__ == "__main__":
-    import rbbt
-
-    filename = "/home/miki/test/numeric.tsv"
-    ds = tsv(filename)
-
-    dl = DataLoader(ds, batch_size=1)
-
-    for f, l in iter(dl):
-        print(".")
-        print(f[0,:])
-        print(l[0])
-
-
-
-

data/python/rbbt_dm/huggingface.py

@@ -1,4 +1,6 @@
 #{{{ LOAD MODEL
+import datasets
+import rbbt
 
 def import_module_class(module, class_name):
     if (not module == None):
@@ -57,12 +59,15 @@ def eval_model(model, tokenizer, texts, return_logits = True):
 #{{{ TRAIN AND PREDICT
 
 def load_tsv(tsv_file):
-    from datasets import load_dataset
-    return load_dataset('csv', data_files=[tsv_file], sep="\t")
+    tsv = rbbt.tsv(tsv_file)
+    print(tsv)
+    ds = datasets.Dataset.from_pandas(tsv)
+    d = datasets.DatasetDict()
+    d["train"] = ds
+    return d
 
 def load_json(json_file):
-    from datasets import load_dataset
-    return load_dataset('json', data_files=[json_file])
+    return datasets.load_dataset('json', data_files=[json_file])
 
 def tokenize_dataset(tokenizer, dataset):
     return dataset.map(lambda subset: subset if ("input_ids" in subset.keys()) else tokenizer(subset["text"], truncation=True), batched=True)

data/python/rbbt_dm/util.py

@@ -28,3 +28,5 @@ def data_directory():
     from pathlib import Path
     print(Path.home())
 
+def model_device(model):
+    return next(model.parameters()).device

data/test/rbbt/vector/model/test_huggingface.rb

@@ -42,7 +42,7 @@ class TestHuggingface < Test::Unit::TestCase
     assert_equal 5, tokenizer.call("This is a sentence that has several words", truncation: true)["input_ids"].__len__
   end
 
-  def _test_sst_eval
+  def test_sst_eval
     TmpFile.with_file do |dir|
       checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 
@@ -56,7 +56,7 @@ class TestHuggingface < Test::Unit::TestCase
   end
 
 
-  def test_sst_train
+  def _test_sst_train
     TmpFile.with_file do |dir|
       checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 

data/test/rbbt/vector/model/test_python.rb

@@ -0,0 +1,31 @@
+require File.expand_path(__FILE__).sub(%r(/test/.*), '/test/test_helper.rb')
+require File.expand_path(__FILE__).sub(%r(.*/test/), '').sub(/test_(.*)\.rb/,'\1')
+
+class TestPythonModel < Test::Unit::TestCase
+  def test_linear
+    model = nil
+
+    TmpFile.with_dir do |dir|
+
+      Misc.in_dir dir do
+        Open.write 'model.py', <<-EOF
+class TestModel:
+  def __init__(self, delta):
+    self.delta = delta
+    
+  def eval(self, x):
+    return [e + self.delta for e in x]
+        EOF
+        model = PythonModel.new dir, 'TestModel', :model, delta: 1
+        
+        assert_equal 2, model.eval(1)
+        assert_equal [4, 6], model.eval_list([3, 5])
+
+        model = PythonModel.new dir, 'TestModel', :model, delta: 2
+        
+        assert_equal 3, model.eval(1)
+      end
+    end
+  end
+end
+

data/test/rbbt/vector/model/test_pytorch_lightning.rb

@@ -2,82 +2,96 @@ require File.join(File.expand_path(File.dirname(__FILE__)), '../../..', 'test_he
 require 'rbbt/vector/model/pytorch_lightning'
 
 class TestPytorchLightning < Test::Unit::TestCase
-  def test_clustering
-    nsamples = 10
-    ngenes = 10000
-    samples = nsamples.times.collect{|i| "Sample-#{i}" }
-    data = TSV.setup({}, :key_field => "Gene", :fields => samples + ["cluster"], :type => :list, :cast => :to_f)
-
-    profiles = []
-    p0 = 3
-    p1 = 7
-    profiles[0] = nsamples.times.collect{ rand() + p0 }
-    profiles[1] = nsamples.times.collect{ rand() + p1 }
-
-    ngenes.times do |genen|
-      gene = "Gene-#{genen}"
-      cluster = genen % 2 
-      values = profiles[cluster].collect do |m|
-        rand() + m
-      end
-      data[gene] = values + [cluster]
-    end
+  def test_regresion
+    points = 10
+    a = 1
+    b = 1
 
+    x = (0..points - 1)
+    y = points.times.collect{|p| p }
+    
     python = <<~EOF
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader
-from torch.utils.data import random_split
-from torchvision.datasets import MNIST
-from torchvision import transforms
 import pytorch_lightning as pl
+import numpy as np
+import torch
+from torch.nn import MSELoss
+from torch.optim import Adam
+from torch.utils.data import DataLoader, Dataset
+import torch.nn as nn
+
+
+class SimpleDataset(Dataset):
+    def __init__(self):
+        X = np.arange(10000)
+        y = X * 2
+        X = [[_] for _ in X]
+        y = [[_] for _ in y]
+        self.X = torch.Tensor(X)
+        self.y = torch.Tensor(y)
+
+    def __len__(self):
+        return len(self.y)
+
+    def __getitem__(self, idx):
+        return {"X": self.X[idx], "y": self.y[idx]}
+
 
 class TestPytorchLightningModel(pl.LightningModule):
-  def __init__(self, input_size=10, internal_dim=1):
-    super().__init__()
-    self.model = nn.Tanh()
-
-  def configure_optimizers(self):
-    optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
-    return optimizer
-
-  @torch.cuda.amp.autocast(True)
-  def forward(self, x):
-    x = x.to(self.dtype)
-    return self.model(x).squeeze()
-
-  @torch.cuda.amp.autocast(True)
-  def training_step(self, train_batch, batch_idx):
-    x, y = train_batch
-    x = x.to(self.dtype)
-    y = y.to(self.dtype)
-    y_hat = self.model(x).squeeze()
-    loss = F.mse_loss(y, y_hat) 
-    self.log('train_loss', loss)
-    return loss
-
-  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float64)
-  def validation_step(self, val_batch, batch_idx):
-    x, y = train_batch
-    y_hat = self.model(x)    
-    loss = F.mse_loss(y, y_hat) 
-    self.log('val_loss', loss)
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Linear(1, 1)
+        self.criterion = MSELoss()
+
+    def forward(self, inputs, labels=None):
+        outputs = self.fc(inputs)
+        loss = 0
+        if labels is not None:
+            loss = self.criterion(outputs, labels)
+        return loss, outputs
+
+    def train_dataloader(self):
+        dataset = SimpleDataset()
+        return DataLoader(dataset, batch_size=1000)
 
+    def training_step(self, batch, batch_idx):
+        input_ids = batch["X"]
+        labels = batch["y"]
+        loss, outputs = self(input_ids, labels)
+        return {"loss": loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=0.1)
+        return optimizer
     EOF
 
-    with_python(python) do |pkg|
-      model = PytorchLightningModel.new pkg , "TestPytorchLightningModel", nil, model_args: {internal_dim: 1}
-      TmpFile.with_file(data.to_s) do |data_file|
-        ds = RbbtPython.call_method "rbbt_dm", :tsv, filename: data_file
-        model.loader = RbbtPython.class_new_obj("torch.utils.data", :DataLoader, dataset: ds, batch_size: 64)
-        model.trainer = RbbtPython.class_new_obj("pytorch_lightning", "Trainer", gpus: 1, max_epochs: 5, precision: 16)
-      end
+    TmpFile.with_dir do |dir|
+      Open.write(File.join(dir, 'model.py'), python)
+      model = PytorchLightningModel.new dir, "TestPytorchLightningModel"
+      model.init
+
+      model.trainer = RbbtPython.class_new_obj("pytorch_lightning", "Trainer", max_epochs: 10, precision: 16)
+      model.init
+
       model.train
-      encoding = model.eval_list(data.values.collect{|v| v[0..-2] }).detach().cpu().numpy()
-      iii encoding[0..10]
+
+      w = model.get_weights('fc').to_ruby.first.first
+
+      assert w > 1.8
+      assert w < 2.2
+
+      res = model.eval(10.0)
+      assert_equal res, (10 * w)
+      assert res > 1.8 * 10.0
+      assert res < 2.2 * 10.0
+
+      res = model.eval([10.0])
+      res = model.eval_list([[10.0], [11.2], [14.3]])
+      assert_equal 3, RbbtPython.numpy2ruby(res).length
+
+      model = VectorModel.new dir
+      model.init
+
     end
   end
-
 end
 

data/test/rbbt/vector/model/test_torch.rb

@@ -0,0 +1,61 @@
+require File.expand_path(__FILE__).sub(%r(/test/.*), '/test/test_helper.rb')
+require File.expand_path(__FILE__).sub(%r(.*/test/), '').sub(/test_(.*)\.rb/,'\1')
+
+class TestTorch < Test::Unit::TestCase
+  def test_linear
+    model = nil
+
+    TmpFile.with_dir do |dir|
+
+      # Create model
+      
+      model = TorchModel.new dir
+      model.model = RbbtPython.torch.nn.Linear.new(1, 1)
+      model.criterion = RbbtPython.torch.nn.MSELoss.new()
+
+      model.extract_features do |f|
+        [f]
+      end
+
+      model.post_process do |v,list|
+        list ? v.to_ruby.collect{|vv| vv.first } :  v.to_ruby.first
+      end
+
+      # Train model
+      
+      model.add 5.0, [10.0]
+      model.add 10.0, [20.0]
+
+      model.training_args[:epochs] = 1000
+      model.train
+
+      w = model.get_weights.to_ruby.first.first
+
+      assert w > 1.8
+      assert w < 2.2
+
+      # Load the model again
+
+      model = VectorModel.new dir
+
+      # Test model
+
+      y = model.eval(100.0)
+
+      assert(y > 150.0)
+      assert(y < 250.0)
+
+      test = [1.0, 5.0, 10.0, 20.0]
+      input_sum = Misc.sum(test)
+      sum = Misc.sum(model.eval_list(test))
+      assert sum > 0.8 * input_sum * 2
+      assert sum < 1.2 * input_sum * 2
+
+      w = TorchModel.get_weights(model.model).to_ruby.first.first
+
+      assert w > 1.8
+      assert w < 2.2
+    end
+  end
+end
+

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rbbt-dm
 version: !ruby/object:Gem::Version
-  version: 1.2.9
+  version: 1.2.10
 platform: ruby
 authors:
 - Miguel Vazquez
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-08-30 00:00:00.000000000 Z
+date: 2023-12-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rbbt-util
@@ -108,12 +108,17 @@ files:
 - lib/rbbt/vector/model.rb
 - lib/rbbt/vector/model/huggingface.rb
 - lib/rbbt/vector/model/huggingface/masked_lm.rb
+- lib/rbbt/vector/model/python.rb
 - lib/rbbt/vector/model/pytorch_lightning.rb
 - lib/rbbt/vector/model/random_forest.rb
 - lib/rbbt/vector/model/spaCy.rb
 - lib/rbbt/vector/model/svm.rb
 - lib/rbbt/vector/model/tensorflow.rb
 - lib/rbbt/vector/model/torch.rb
+- lib/rbbt/vector/model/torch/dataloader.rb
+- lib/rbbt/vector/model/torch/helpers.rb
+- lib/rbbt/vector/model/torch/introspection.rb
+- lib/rbbt/vector/model/torch/load_and_save.rb
 - lib/rbbt/vector/model/util.rb
 - python/rbbt_dm/__init__.py
 - python/rbbt_dm/atcold/__init__.py
@@ -143,10 +148,12 @@ files:
 - test/rbbt/test_stan.rb
 - test/rbbt/vector/model/huggingface/test_masked_lm.rb
 - test/rbbt/vector/model/test_huggingface.rb
+- test/rbbt/vector/model/test_python.rb
 - test/rbbt/vector/model/test_pytorch_lightning.rb
 - test/rbbt/vector/model/test_spaCy.rb
 - test/rbbt/vector/model/test_svm.rb
 - test/rbbt/vector/model/test_tensorflow.rb
+- test/rbbt/vector/model/test_torch.rb
 - test/rbbt/vector/test_model.rb
 - test/test_helper.rb
 homepage: http://github.com/mikisvaz/rbbt-phgx
@@ -167,7 +174,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.19
+rubygems_version: 3.5.0.dev
 signing_key: 
 specification_version: 4
 summary: Data-mining and statistics
@@ -182,9 +189,11 @@ test_files:
 - test/rbbt/test_stan.rb
 - test/rbbt/vector/model/huggingface/test_masked_lm.rb
 - test/rbbt/vector/model/test_huggingface.rb
+- test/rbbt/vector/model/test_python.rb
 - test/rbbt/vector/model/test_pytorch_lightning.rb
 - test/rbbt/vector/model/test_spaCy.rb
 - test/rbbt/vector/model/test_svm.rb
 - test/rbbt/vector/model/test_tensorflow.rb
+- test/rbbt/vector/model/test_torch.rb
 - test/rbbt/vector/test_model.rb
 - test/test_helper.rb