rbbt-dm 1.2.3 → 1.2.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2ff72107967b0f7c654697f3a7b3c0ef10f7a5264d775117f12f74084a2819b2
-  data.tar.gz: 6b9a58b5a2723c095332f79a37d9c1c7f4bc1431410f23a55beeed1c3b52f7ad
+  metadata.gz: 9744ab9faeaf4f9cc04947eb11103dbf0694dda624f805a5c6be27bb22af81ce
+  data.tar.gz: d3a3903aa276a69e20cbd71213286449db396ecf5f6a4b4d80a64ab299041fbb
 SHA512:
-  metadata.gz: a0fb4198cb0be3aa5253df0f655ee230621dd26a31956a774fffe95eac35f4c8b558a41c0e340c25c6eef463760ff6230b967f09eb09671b2078a50066067384
-  data.tar.gz: 0bd6c3667a8ec26ed092c54e78176671807ed0634e136e497303e34a17b2740e5e023041bc06389ac187de39b54942b9b1c5cd77abbc067c89250424654b6974
+  metadata.gz: 263fb609b37522874426bcd79374760399b4a9aaab443ae6d74c727f2d148474dd71ee0b2cfda7a50131dafbc314f66352f4285562a75f62144d2e05ccd214c7
+  data.tar.gz: 1e0426429a38028a19b3f8c955e975138199c791dad8691de7fb760a5cbec3304f19341906a4457563a90de965ee2f12a5a639b928866b46258af2507eeb39fa

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

@@ -0,0 +1,160 @@
+require 'rbbt/vector/model'
+require 'rbbt/util/python'
+
+RbbtPython.add_path Rbbt.python.find(:lib)
+RbbtPython.init_rbbt
+
+class HuggingfaceModel < VectorModel
+
+  attr_accessor :checkpoint, :task, :locate_tokens, :class_labels, :class_weights, :training_args
+
+  def self.tsv_dataset(tsv_dataset_file, elements, labels = nil)
+
+    if labels
+      Open.write(tsv_dataset_file) do |ffile|
+        ffile.puts ["label", "text"].flatten * "\t"
+        elements.zip(labels).each do |element,label|
+          ffile.puts [label, element].flatten * "\t"
+        end
+      end
+    else
+      Open.write(tsv_dataset_file) do |ffile|
+        ffile.puts ["text"].flatten * "\t"
+        elements.each{|element| ffile.puts element }
+      end
+    end
+
+    tsv_dataset_file
+  end
+
+  def self.call_method(name, *args)
+    RbbtPython.import_method("rbbt_dm.huggingface", name).call(*args)
+  end
+
+  def call_method(name, *args)
+    HuggingfaceModel.call_method(name, *args)
+  end
+
+  #def input_tsv_file
+  #  File.join(@directory, 'dataset.tsv') if @directory
+  #end
+
+  #def checkpoint_dir
+  #  File.join(@directory, 'checkpoints') if @directory
+  #end
+
+  def self.run_model(model, tokenizer, elements, labels = nil, training_args = {}, class_weights = nil)
+    TmpFile.with_file do |tmpfile|
+      tsv_file = File.join(tmpfile, 'dataset.tsv')
+
+      if training_args
+        training_args = training_args.dup
+        checkpoint_dir = training_args.delete(:checkpoint_dir)
+      end
+
+      checkpoint_dir = File.join(tmpfile, 'checkpoints')
+
+      Open.mkdir File.dirname(tsv_file)
+      Open.mkdir File.dirname(checkpoint_dir)
+
+      if labels
+        training_args_obj = call_method(:training_args, checkpoint_dir, **training_args)
+        call_method(:train_model, model, tokenizer, training_args_obj, tsv_dataset(tsv_file, elements, labels), class_weights)
+      else
+        locate_tokens, training_args = training_args, {}
+        if Array === elements
+          training_args_obj = call_method(:training_args, checkpoint_dir)
+          call_method(:predict_model, model, tokenizer, training_args_obj, tsv_dataset(tsv_file, elements), locate_tokens)
+        else
+          call_method(:eval_model, model, tokenizer, [elements], locate_tokens)
+        end
+      end
+    end
+  end
+
+  def init_model
+    @model, @tokenizer = call_method(:load_model_and_tokenizer, @task, @checkpoint)
+  end
+
+  def reset_model
+    init_model
+  end
+
+  def initialize(task, initial_checkpoint = nil, *args)
+    super(*args)
+    @task = task
+
+    @checkpoint = model_file && File.exists?(model_file)? model_file : initial_checkpoint
+
+    init_model
+
+    @locate_tokens = @tokenizer.special_tokens_map["mask_token"]  if @task == "MaskedLM"
+
+    @training_args = {}
+
+    train_model do |file,elements,labels|
+      HuggingfaceModel.run_model(@model, @tokenizer, elements, labels, @training_args, @class_weights)
+
+      @model.save_pretrained(file) if file
+      @tokenizer.save_pretrained(file) if file
+    end
+
+    eval_model do |file,elements|
+      @model, @tokenizer = HuggingfaceModel.call_method(:load_model_and_tokenizer, @task, @checkpoint)
+      HuggingfaceModel.run_model(@model, @tokenizer, elements, nil, @locate_tokens)
+    end
+
+    post_process do |result|
+      if result.respond_to?(:predictions)
+        single = false
+        predictions = result.predictions
+      elsif result["token_positions"]
+        predictions = result["result"].predictions
+        token_positions = result["token_positions"]
+      else
+        single = true
+        predictions = result["logits"]
+      end
+
+      result = case @task
+               when "SequenceClassification"
+                 RbbtPython.collect(predictions) do |logits|
+                   logits = RbbtPython.numpy2ruby logits
+                   best_class = logits.index logits.max
+                   best_class = @class_labels[best_class] if @class_labels
+                   best_class
+                 end
+               when "MaskedLM"
+                 all_token_positions = token_positions.to_a
+
+                 i = 0
+                 RbbtPython.collect(predictions) do |item_logits|
+                   item_token_positions = all_token_positions[i]
+                   i += 1
+
+                   item_logits = RbbtPython.numpy2ruby(item_logits)
+                   item_masks = item_token_positions.collect do |token_positions|
+
+                     best = item_logits.values_at(*token_positions).collect do |logits|
+                       best_token, best_score = nil
+                       logits.each_with_index do |v,i|
+                         if best_score.nil? || v > best_score
+                           best_token, best_score = i, v
+                         end
+                       end
+                       best_token
+                     end
+
+                     best.collect{|b| @tokenizer.decode(b) } * "|"
+                   end
+                   Array === @locate_tokens ? item_masks : item_masks.first
+                 end
+               else
+                 logits
+               end
+
+      single ? result.first : result
+    end
+  end
+end
+

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

@@ -1,13 +1,12 @@
 require 'rbbt/vector/model'
 require 'rbbt/util/python'
 
+RbbtPython.add_path Rbbt.python.find(:lib)
 RbbtPython.init_rbbt
 
 class HuggingfaceModel < VectorModel
 
-  attr_accessor :checkpoint, :task, :locate_tokens, :class_labels
-
-  def tsv_dataset(tsv_dataset_file, elements, labels = nil)
+  def self.tsv_dataset(tsv_dataset_file, elements, labels = nil)
 
     if labels
       Open.write(tsv_dataset_file) do |ffile|
@@ -26,59 +25,74 @@ class HuggingfaceModel < VectorModel
     tsv_dataset_file
   end
 
-  def call_method(name, *args)
-    RbbtPython.import_method("rbbt_dm.huggingface", name).call(*args)
-  end
-
-  def input_tsv_file
-    File.join(@directory, 'dataset.tsv') if @directory
-  end
+  def initialize(task, checkpoint, *args)
+    options = args.pop if Hash === args.last
+    options = Misc.add_defaults options, :task => task, :checkpoint => checkpoint
+    super(*args)
+    @model_options ||= {}
+    @model_options.merge!(options)
 
-  def checkpoint_dir
-    File.join(@directory, 'checkpoints') if @directory
-  end
+    eval_model do |directory,texts|
+      checkpoint = directory && File.directory?(directory) ? directory : @model_options[:checkpoint]
 
-  def run_model(elements, labels = nil)
-    TmpFile.with_file do |tmpfile|
-      tsv_file = input_tsv_file || File.join(tmpfile, 'dataset.tsv')
-      output_dir = checkpoint_dir || File.join(tmpfile, 'checkpoints')
+      if @model.nil?
+        @model, @tokenizer = RbbtPython.call_method("rbbt_dm.huggingface", :load_model_and_tokenizer, @model_options[:task], checkpoint)
+      end
+      
+      if Array === texts
 
-      Open.mkdir File.dirname(output_dir)
-      Open.mkdir File.dirname(tsv_file)
+        if @model_options.include?(:locate_tokens)
+          locate_tokens = @model_options[:locate_tokens]
+        elsif @model_options[:task] == "MaskedLM"
+          @model_options[:locate_tokens] = locate_tokens = @tokenizer.special_tokens_map["mask_token"] 
+        end
 
-      if labels
-        training_args = call_method(:training_args, output_dir)
-        call_method(:train_model, @model, @tokenizer, training_args, tsv_dataset(tsv_file, elements, labels))
-      else
-        if Array === elements
-          training_args = call_method(:training_args, output_dir)
-          call_method(:predict_model, @model, @tokenizer, training_args, tsv_dataset(tsv_file, elements), @locate_tokens)
+        if @directory
+          tsv_file = File.join(@directory, 'dataset.tsv')
+          checkpoint_dir = File.join(@directory, 'checkpoints')
         else
-          call_method(:eval_model, @model, @tokenizer, [elements], @locate_tokens)
+          tmpdir = TmpFile.tmp_file
+          Open.mkdir tmpdir
+          tsv_file = File.join(tmpdir, 'dataset.tsv')
+          checkpoint_dir = File.join(tmpdir, 'checkpoints')
+        end
+
+        dataset_file = HuggingfaceModel.tsv_dataset(tsv_file, texts)
+        training_args_obj = RbbtPython.call_method("rbbt_dm.huggingface", :training_args, checkpoint_dir, @model_options[:training_args])
+
+        begin
+          RbbtPython.call_method("rbbt_dm.huggingface", :predict_model, @model, @tokenizer, training_args_obj, dataset_file, locate_tokens)
+        ensure
+          Open.rm_rf tmpdir if tmpdir
         end
+      else
+        RbbtPython.call_method("rbbt_dm.huggingface", :eval_model, @model, @tokenizer, [texts], locate_tokens)
       end
     end
-  end
-
-  def initialize(task, initial_checkpoint = nil, *args)
-    super(*args)
-    @task = task
 
-    @checkpoint = model_file && File.exists?(model_file)? model_file : initial_checkpoint
+    train_model do |directory,texts,labels|
+      checkpoint = directory && File.directory?(directory) ? directory : @model_options[:checkpoint]
+      @model, @tokenizer = RbbtPython.call_method("rbbt_dm.huggingface", :load_model_and_tokenizer, @model_options[:task], checkpoint)
 
-    @model, @tokenizer = call_method(:load_model_and_tokenizer, @task, @checkpoint)
+      if @directory
+        tsv_file = File.join(@directory, 'dataset.tsv')
+        checkpoint_dir = File.join(@directory, 'checkpoints')
+      else
+        tmpdir = TmpFile.tmp_file
+        Open.mkdir tmpdir
+        tsv_file = File.join(tmpdir, 'dataset.tsv')
+        checkpoint_dir = File.join(tmpdir, 'checkpoints')
+      end
 
-    @locate_tokens = @tokenizer.special_tokens_map["mask_token"]  if @task == "MaskedLM"
+      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)
 
-    train_model do |file,elements,labels|
-      run_model(elements, labels)
+      RbbtPython.call_method("rbbt_dm.huggingface", :train_model, @model, @tokenizer, training_args_obj, dataset_file, @model_options[:class_weights])
 
-      @model.save_pretrained(file) if file
-      @tokenizer.save_pretrained(file) if file
-    end
+      Open.rm_rf tmpdir if tmpdir
 
-    eval_model do |file,elements|
-      run_model(elements)
+      @model.save_pretrained(directory) if directory
+      @tokenizer.save_pretrained(directory) if directory
     end
 
     post_process do |result|
@@ -93,12 +107,13 @@ class HuggingfaceModel < VectorModel
         predictions = result["logits"]
       end
 
-      result = case @task
+      task, class_labels, locate_tokens = @model_options.values_at :task, :class_labels, :locate_tokens
+      result = case task
                when "SequenceClassification"
                  RbbtPython.collect(predictions) do |logits|
                    logits = RbbtPython.numpy2ruby logits
                    best_class = logits.index logits.max
-                   best_class = @class_labels[best_class] if @class_labels
+                   best_class = class_labels[best_class] if class_labels
                    best_class
                  end
                when "MaskedLM"
@@ -124,7 +139,7 @@ class HuggingfaceModel < VectorModel
 
                      best.collect{|b| @tokenizer.decode(b) } * "|"
                    end
-                   Array === @locate_tokens ? item_masks : item_masks.first
+                   Array === locate_tokens ? item_masks : item_masks.first
                  end
                else
                  logits
@@ -132,9 +147,15 @@ class HuggingfaceModel < VectorModel
 
       single ? result.first : result
     end
+
+
+    save_models if @directory
   end
-end
 
-if __FILE__ == $0
+  def reset_model
+    @model, @tokenizer = nil
+    Open.rm @model_file
+  end
 
 end
+

data/lib/rbbt/vector/model.rb

@@ -4,6 +4,7 @@ require 'rbbt/vector/model/util'
 class VectorModel
   attr_accessor :directory, :model_file, :extract_features, :train_model, :eval_model, :post_process
   attr_accessor :features, :names, :labels, :factor_levels
+  attr_accessor :model_options
 
   def extract_features(&block)
     @extract_features = block if block_given?
@@ -126,7 +127,7 @@ cat(paste(label, sep="\\n", collapse="\\n"));
     instance_eval code, file
   end
 
-  def initialize(directory = nil, extract_features = nil, train_model = nil, eval_model = nil, names = nil, factor_levels = nil)
+  def initialize(directory = nil, extract_features = nil, train_model = nil, eval_model = nil, post_process = nil, names = nil, factor_levels = nil)
     @directory = directory
     if @directory
       FileUtils.mkdir_p @directory unless File.exists?(@directory)
@@ -135,10 +136,18 @@ cat(paste(label, sep="\\n", collapse="\\n"));
       @extract_features_file = File.join(@directory, "features")
       @train_model_file = File.join(@directory, "train_model")
       @eval_model_file = File.join(@directory, "eval_model")
+      @post_process_file = File.join(@directory, "post_process")
       @train_model_file_R = File.join(@directory, "train_model.R")
       @eval_model_file_R = File.join(@directory, "eval_model.R")
+      @post_process_file_R = File.join(@directory, "post_process.R")
       @names_file = File.join(@directory, "feature_names")
       @levels_file = File.join(@directory, "levels")
+      @options_file = File.join(@directory, "options.json")
+
+      if File.exists?(@options_file)
+        @model_options = JSON.parse(Open.read(@options_file))
+        IndiferentHash.setup(@model_options)
+      end
     end
 
     if extract_features.nil?
@@ -169,6 +178,17 @@ cat(paste(label, sep="\\n", collapse="\\n"));
       @eval_model = eval_model
     end
 
+    if post_process.nil?
+      if @post_process_file && File.exists?(@post_process_file)
+        @post_process = __load_method @post_process_file
+      elsif @post_process_file_R && File.exists?(@post_process_file_R)
+        @post_process = Open.read(@post_process_file_R)
+      end
+    else
+      @post_process = post_process
+    end
+
+
     if names.nil?
       if @names_file && File.exists?(@names_file)
         @names = Open.read(@names_file).split("\n")
@@ -240,8 +260,20 @@ cat(paste(label, sep="\\n", collapse="\\n"));
       Open.write(@eval_model_file_R, eval_model)
     end
 
+    case 
+    when Proc === post_process
+      begin
+        Open.write(@post_process_file, post_process.source)
+      rescue
+      end
+    when String === post_process
+      Open.write(@post_process_file_R, post_process)
+    end
+
+
     Open.write(@levels_file, @factor_levels.to_yaml) if @factor_levels
     Open.write(@names_file, @names * "\n" + "\n") if @names
+    Open.write(@options_file, @model_options.to_json) if @model_options
   end
 
   def train
@@ -251,7 +283,7 @@ cat(paste(label, sep="\\n", collapse="\\n"));
     when String === @train_model
       VectorModel.R_train(@model_file,  @features, @labels, train_model, @names, @factor_levels)
     end
-    save_models
+    save_models if @directory
   end
 
   def run(code)
@@ -299,38 +331,6 @@ cat(paste(label, sep="\\n", collapse="\\n"));
     result
   end
 
-  #def cross_validation(folds = 10)
-  #  saved_features = @features
-  #  saved_labels = @labels
-  #  seq = (0..features.length - 1).to_a
-
-  #  chunk_size = features.length / folds
-
-  #  acc = []
-  #  folds.times do
-  #    seq = seq.shuffle
-  #    eval_chunk = seq[0..chunk_size]
-  #    train_chunk = seq[chunk_size.. -1]
-
-  #    eval_features = @features.values_at *eval_chunk
-  #    eval_labels = @labels.values_at *eval_chunk
-
-  #    @features = @features.values_at *train_chunk
-  #    @labels = @labels.values_at *train_chunk
-
-  #    train
-  #    predictions = eval_list eval_features, false
-
-  #    acc << predictions.zip(eval_labels).collect{|pred,lab| pred - lab < 0.5 ? 1 : 0}.inject(0){|acc,e| acc +=e} / chunk_size
-
-  #    @features = saved_features
-  #    @labels = saved_labels
-  #  end
-
-  #  acc
-  #end
-  #
-  
   def self.f1_metrics(test, predicted, good_label = nil)
     tp, tn, fp, fn, pr, re, f1 = [0, 0, 0, 0, nil, nil, nil]
 
@@ -413,6 +413,7 @@ cat(paste(label, sep="\\n", collapse="\\n"));
         @features = train_set
         @labels = train_labels
 
+        self.reset_model if self.respond_to? :reset_model
         self.train
         predictions = self.eval_list test_set, false
 
@@ -437,6 +438,7 @@ cat(paste(label, sep="\\n", collapse="\\n"));
       @features = orig_features
       @labels = orig_labels
     end unless folds == -1
+    self.reset_model if self.respond_to? :reset_model
     self.train unless folds == 1
     res
   end

data/python/rbbt_dm/__init__.py

@@ -0,0 +1 @@
+# Keep

data/python/rbbt_dm/huggingface.py

@@ -17,6 +17,17 @@ def load_model_and_tokenizer(task, checkpoint):
     tokenizer = load_tokenizer(task, checkpoint)
     return model, tokenizer
 
+def load_model_and_tokenizer_from_directory(directory):
+    import os
+    import json
+    options_file = os.path.join(directory, 'options.json')
+    f = open(options_file, "r")
+    options = json.load(f.read())
+    f.close()
+    task = options["task"]
+    checkpoint = options["checkpoint"]
+    return load_model_and_tokenizer(task, checkpoint)
+
 #{{{ SIMPLE EVALUATE
 
 def forward(model, features):
@@ -51,17 +62,40 @@ def training_args(*args, **kwargs):
     return training_args
 
 
-def train_model(model, tokenizer, training_args, tsv_file):
+def train_model(model, tokenizer, training_args, tsv_file, class_weights=None):
     from transformers import Trainer
 
     tokenized_dataset = tsv_dataset(tokenizer, tsv_file)
 
-    trainer = Trainer(
-            model,
-            training_args,
-            train_dataset = tokenized_dataset["train"],
-            tokenizer = tokenizer
-            )
+    if (not class_weights == None):
+        import torch
+        from torch import nn
+
+        class WeightTrainer(Trainer):
+            def compute_loss(self, model, inputs, return_outputs=False):
+                labels = inputs.get("labels")
+                # forward pass
+                outputs = model(**inputs)
+                logits = outputs.get('logits')
+                # compute custom loss
+                loss_fct = nn.CrossEntropyLoss(weight=torch.tensor(class_weights).to(model.device))
+                loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
+                return (loss, outputs) if return_outputs else loss
+
+        trainer = WeightTrainer(
+                model,
+                training_args,
+                train_dataset = tokenized_dataset["train"],
+                tokenizer = tokenizer
+                )
+    else:
+
+        trainer = Trainer(
+                model,
+                training_args,
+                train_dataset = tokenized_dataset["train"],
+                tokenizer = tokenizer
+                )
 
     trainer.train()
 
@@ -90,7 +124,6 @@ def find_tokens_in_input(dataset, token_ids):
     return position_rows
 
 
-
 def predict_model(model, tokenizer, training_args, tsv_file, locate_tokens = None):
     from transformers import Trainer
 
@@ -110,3 +143,4 @@ def predict_model(model, tokenizer, training_args, tsv_file, locate_tokens = Non
     else:
         return result
 
+

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

@@ -3,16 +3,46 @@ require 'rbbt/vector/model/huggingface'
 
 class TestHuggingface < Test::Unit::TestCase
 
+  def test_options
+    TmpFile.with_file do |dir|
+      checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
+      task = "SequenceClassification"
+
+      model = HuggingfaceModel.new task, checkpoint, dir, :class_labels => %w(bad good)
+      iii model.eval "This is dog"
+      iii model.eval "This is cat"
+      iii model.eval(["This is dog", "This is cat"])
+
+      model = VectorModel.new dir
+      iii model.eval(["This is dog", "This is cat"])
+    end
+  end
+
+  def test_pipeline
+    require 'rbbt/util/python'
+    model = VectorModel.new
+    model.post_process do |elements|
+      elements.collect{|e| e['label'] }
+    end
+    model.eval_model do |file, elements|
+      RbbtPython.run :transformers do 
+        classifier ||= transformers.pipeline("sentiment-analysis")
+        classifier.call(elements)
+      end
+    end
+
+    assert_equal ["POSITIVE"], model.eval("I've been waiting for a HuggingFace course my whole life.")
+  end
+
   def test_sst_eval
     TmpFile.with_file do |dir|
       checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 
       model = HuggingfaceModel.new "SequenceClassification", checkpoint, dir
 
-      model.class_labels = ["Bad", "Good"]
+      model.model_options[:class_labels] = ["Bad", "Good"]
 
       assert_equal ["Bad", "Good"], model.eval(["This is dog", "This is cat"])
-
     end
   end
 
@@ -22,7 +52,8 @@ class TestHuggingface < Test::Unit::TestCase
       checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 
       model = HuggingfaceModel.new "SequenceClassification", checkpoint, dir
-      model.class_labels = ["Bad", "Good"]
+
+      model.model_options[:class_labels] = %w(Bad Good)
 
       assert_equal ["Bad", "Good"], model.eval(["This is dog", "This is cat"])
 
@@ -33,6 +64,9 @@ class TestHuggingface < Test::Unit::TestCase
       model.train
 
       assert_equal ["Good", "Good"], model.eval(["This is dog", "This is cat"])
+
+      model = VectorModel.new dir
+      assert_equal ["Good", "Good"], model.eval(["This is dog", "This is cat"])
     end
   end
 
@@ -40,7 +74,7 @@ class TestHuggingface < Test::Unit::TestCase
     checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 
     model = HuggingfaceModel.new "SequenceClassification", checkpoint
-    model.class_labels = ["Bad", "Good"]
+    model.model_options[:class_labels] = ["Bad", "Good"]
 
     assert_equal ["Bad", "Good"], model.eval(["This is dog", "This is cat"])
 
@@ -58,7 +92,7 @@ class TestHuggingface < Test::Unit::TestCase
       checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 
       model = HuggingfaceModel.new "SequenceClassification", checkpoint, dir
-      model.class_labels = ["Bad", "Good"]
+      model.model_options[:class_labels] = ["Bad", "Good"]
 
       assert_equal ["Bad", "Good"], model.eval(["This is dog", "This is cat"])
 
@@ -69,15 +103,20 @@ class TestHuggingface < Test::Unit::TestCase
       model.train
 
       model = HuggingfaceModel.new "SequenceClassification", checkpoint, dir
-      model.class_labels = ["Bad", "Good"]
 
       assert_equal ["Good", "Good"], model.eval(["This is dog", "This is cat"])
 
-      model = HuggingfaceModel.new "SequenceClassification", model.model_file
-      model.class_labels = ["Bad", "Good"]
+      model_file = model.model_file
+
+      model = HuggingfaceModel.new "SequenceClassification", model_file
+      model.model_options[:class_labels] = ["Bad", "Good"]
 
       assert_equal ["Good", "Good"], model.eval(["This is dog", "This is cat"])
 
+      model = VectorModel.new dir
+
+      assert_equal "Good", model.eval("This is dog")
+
     end
   end
 
@@ -108,9 +147,7 @@ class TestHuggingface < Test::Unit::TestCase
     model = HuggingfaceModel.new "MaskedLM", checkpoint
     assert_equal 3, model.eval(["Paris is the [MASK] of the France.", "The [MASK] worked very hard all the time.", "The [MASK] arrested the dangerous [MASK]."]).
       reject{|v| v.empty?}.length
-
   end
 
-
 end
 

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rbbt-dm
 version: !ruby/object:Gem::Version
-  version: 1.2.3
+  version: 1.2.6
 platform: ruby
 authors:
 - Miguel Vazquez
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-02-04 00:00:00.000000000 Z
+date: 2023-02-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rbbt-util
@@ -107,12 +107,14 @@ files:
 - lib/rbbt/statistics/rank_product.rb
 - lib/rbbt/tensorflow.rb
 - lib/rbbt/vector/model.rb
+- lib/rbbt/vector/model/huggingface.old.rb
 - lib/rbbt/vector/model/huggingface.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/util.rb
+- python/rbbt_dm/__init__.py
 - python/rbbt_dm/huggingface.py
 - share/R/MA.R
 - share/R/barcode.R