Replaced task classes with pipeline method and added Model class

.github/workflows/build.yml

@@ -11,17 +11,6 @@ jobs:
           bundler-cache: true
       - uses: actions/cache@v4
         with:
-          path: models
-          key: models-v3
-        id: cache-models
-      - name: Download models
-        if: steps.cache-models.outputs.cache-hit != 'true'
-        run: |
-          mkdir models
-          cd models
-          wget -q https://github.com/ankane/informers/releases/download/v0.1.0/sentiment-analysis.onnx
-          wget -q https://github.com/ankane/informers/releases/download/v0.1.0/question-answering.onnx
-          wget -q https://github.com/ankane/informers/releases/download/v0.1.0/feature-extraction.onnx
+          path: ~/.cache/informers
+          key: informers
       - run: bundle exec rake test
-        env:
-          MODELS_PATH: models

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## 1.0.0 (unreleased)
 
+- Replaced task classes with `pipeline` method
+- Added `Model` class
 - Dropped support for Ruby < 3.1
 
 ## 0.2.0 (2022-09-06)

README.md

@@ -1,13 +1,8 @@
 # Informers
 
-:slightly_smiling_face: State-of-the-art natural language processing for Ruby
+:fire: Fast [transformer](https://github.com/xenova/transformers.js) inference for Ruby
 
-Supports:
-
-- Sentiment analysis
-- Question answering
-- Named-entity recognition
-- Text generation
+For non-ONNX models, check out [Transformers.rb](https://github.com/ankane/transformers-ruby)
 
 [![Build Status](https://github.com/ankane/informers/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/informers/actions)
 
@@ -21,142 +16,111 @@ gem "informers"
 
 ## Getting Started
 
-- [Sentiment analysis](#sentiment-analysis)
-- [Question answering](#question-answering)
-- [Named-entity recognition](#named-entity-recognition)
-- [Text generation](#text-generation)
-- [Feature extraction](#feature-extraction)
-- [Fill mask](#fill-mask)
+- [Models](#models)
+- [Pipelines](#pipelines)
 
-### Sentiment Analysis
+## Models
 
-First, download the [pretrained model](https://github.com/ankane/informers/releases/download/v0.1.0/sentiment-analysis.onnx).
+### sentence-transformers/all-MiniLM-L6-v2
 
-Predict sentiment
+[Docs](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
 
 ```ruby
-model = Informers::SentimentAnalysis.new("sentiment-analysis.onnx")
-model.predict("This is super cool")
-```
+sentences = ["This is an example sentence", "Each sentence is converted"]
 
-This returns
-
-```ruby
-{label: "positive", score: 0.999855186578301}
+model = Informers::Model.new("sentence-transformers/all-MiniLM-L6-v2")
+embeddings = model.embed(sentences)
 ```
 
-Predict multiple at once
+For a quantized version, use:
 
 ```ruby
-model.predict(["This is super cool", "I didn't like it"])
+model = Informers::Model.new("Xenova/all-MiniLM-L6-v2", quantized: true)
 ```
 
-### Question Answering
-
-First, download the [pretrained model](https://github.com/ankane/informers/releases/download/v0.1.0/question-answering.onnx).
+### Xenova/multi-qa-MiniLM-L6-cos-v1
 
-Ask a question with some context
+[Docs](https://huggingface.co/Xenova/multi-qa-MiniLM-L6-cos-v1)
 
 ```ruby
-model = Informers::QuestionAnswering.new("question-answering.onnx")
-model.predict(
-  question: "Who invented Ruby?",
-  context: "Ruby is a programming language created by Matz"
-)
+query = "How many people live in London?"
+docs = ["Around 9 Million people live in London", "London is known for its financial district"]
+
+model = Informers::Model.new("Xenova/multi-qa-MiniLM-L6-cos-v1")
+query_embedding = model.embed(query)
+doc_embeddings = model.embed(docs)
+scores = doc_embeddings.map { |e| e.zip(query_embedding).sum { |d, q| d * q } }
+doc_score_pairs = docs.zip(scores).sort_by { |d, s| -s }
 ```
 
-This returns
+### mixedbread-ai/mxbai-embed-large-v1
 
-```ruby
-{answer: "Matz", score: 0.9980658360049758, start: 42, end: 46}
-```
-
-Note: The question and context combined are limited to 384 tokens
-
-### Named-Entity Recognition
-
-First, export the [pretrained model](tools/export.md).
-
-Get entities
+[Docs](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1)
 
 ```ruby
-model = Informers::NER.new("ner.onnx")
-model.predict("Nat works at GitHub in San Francisco")
-```
-
-This returns
-
-```ruby
-[
-  {text: "Nat",           tag: "person",   score: 0.9840519576513487, start: 0,  end: 3},
-  {text: "GitHub",        tag: "org",      score: 0.9426134775785775, start: 13, end: 19},
-  {text: "San Francisco", tag: "location", score: 0.9952414982243061, start: 23, end: 36}
+def transform_query(query)
+  "Represent this sentence for searching relevant passages: #{query}"
+end
+
+docs = [
+  transform_query("puppy"),
+  "The dog is barking",
+  "The cat is purring"
 ]
-```
 
-### Text Generation
+model = Informers::Model.new("mixedbread-ai/mxbai-embed-large-v1")
+embeddings = model.embed(docs)
+```
 
-First, export the [pretrained model](tools/export.md).
+## Pipelines
 
-Pass a prompt
+Named-entity recognition
 
 ```ruby
-model = Informers::TextGeneration.new("text-generation.onnx")
-model.predict("As far as I am concerned, I will", max_length: 50)
+ner = Informers.pipeline("ner")
+ner.("Ruby is a programming language created by Matz")
 ```
 
-This returns
+Sentiment analysis
 
-```text
-As far as I am concerned, I will be the first to admit that I am not a fan of the idea of a "free market." I think that the idea of a free market is a bit of a stretch. I think that the idea
+```ruby
+classifier = Informers.pipeline("sentiment-analysis")
+classifier.("We are very happy to show you the 🤗 Transformers library.")
 ```
 
-### Feature Extraction
-
-First, export a [pretrained model](tools/export.md).
+Question answering
 
 ```ruby
-model = Informers::FeatureExtraction.new("feature-extraction.onnx")
-model.predict("This is super cool")
+qa = Informers.pipeline("question-answering")
+qa.("Who invented Ruby?", "Ruby is a programming language created by Matz")
 ```
 
-### Fill Mask
-
-First, export a [pretrained model](tools/export.md).
+Feature extraction
 
 ```ruby
-model = Informers::FillMask.new("fill-mask.onnx")
-model.predict("This is a great <mask>")
+extractor = Informers.pipeline("feature-extraction")
+extractor.("We are very happy to show you the 🤗 Transformers library.")
 ```
 
-## Models
-
-Task | Description | Contributor | License | Link
---- | --- | --- | --- | ---
-Sentiment analysis | DistilBERT fine-tuned on SST-2 | Hugging Face | Apache-2.0 | [Link](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english)
-Question answering | DistilBERT fine-tuned on SQuAD | Hugging Face | Apache-2.0 | [Link](https://huggingface.co/distilbert-base-cased-distilled-squad)
-Named-entity recognition | BERT fine-tuned on CoNLL03 | Bayerische Staatsbibliothek | In-progress | [Link](https://huggingface.co/dbmdz/bert-large-cased-finetuned-conll03-english)
-Text generation | GPT-2 | OpenAI | [Custom](https://github.com/openai/gpt-2/blob/master/LICENSE) | [Link](https://huggingface.co/gpt2)
-
-Some models are [quantized](https://medium.com/microsoftazure/faster-and-smaller-quantized-nlp-with-hugging-face-and-onnx-runtime-ec5525473bb7) to make them faster and smaller.
-
-## Deployment
+## Credits
 
-Check out [Trove](https://github.com/ankane/trove) for deploying models.
+This library was ported from [Transformers.js](https://github.com/xenova/transformers.js) and is available under the same license.
 
-```sh
-trove push sentiment-analysis.onnx
-```
+## Upgrading
 
-## Credits
+### 1.0
 
-This project uses many state-of-the-art technologies:
+Task classes have been replaced with the `pipeline` method.
 
-- [Transformers](https://github.com/huggingface/transformers) for transformer models
-- [Bling Fire](https://github.com/microsoft/BlingFire) and [BERT](https://github.com/google-research/bert) for high-performance text tokenization
-- [ONNX Runtime](https://github.com/Microsoft/onnxruntime) for high-performance inference
+```ruby
+# before
+model = Informers::SentimentAnalysis.new("sentiment-analysis.onnx")
+model.predict("This is super cool")
 
-Some code was ported from Transformers and is available under the same license.
+# after
+model = Informers.pipeline("sentiment-analysis")
+model.("This is super cool")
+```
 
 ## History
 
@@ -177,7 +141,5 @@ To get started with development:
 git clone https://github.com/ankane/informers.git
 cd informers
 bundle install
-
-export MODELS_PATH=path/to/onnx/models
 bundle exec rake test
 ```

informers.gemspec

@@ -3,19 +3,18 @@ require_relative "lib/informers/version"
 Gem::Specification.new do |spec|
   spec.name          = "informers"
   spec.version       = Informers::VERSION
-  spec.summary       = "State-of-the-art natural language processing for Ruby"
+  spec.summary       = "Fast transformer inference for Ruby"
   spec.homepage      = "https://github.com/ankane/informers"
   spec.license       = "Apache-2.0"
 
   spec.author        = "Andrew Kane"
   spec.email         = "[email protected]"
 
-  spec.files         = Dir["*.{md,txt}", "{lib,vendor}/**/*"]
+  spec.files         = Dir["*.{md,txt}", "{lib}/**/*"]
   spec.require_path  = "lib"
 
   spec.required_ruby_version = ">= 3.1"
 
-  spec.add_dependency "blingfire", ">= 0.1.7"
-  spec.add_dependency "numo-narray"
-  spec.add_dependency "onnxruntime", ">= 0.5.1"
+  spec.add_dependency "onnxruntime", ">= 0.9"
+  spec.add_dependency "tokenizers", ">= 0.5.2"
 end

lib/informers.rb

@@ -1,13 +1,32 @@
 # dependencies
-require "blingfire"
-require "numo/narray"
 require "onnxruntime"
+require "tokenizers"
+
+# stdlib
+require "io/console"
+require "json"
+require "open-uri"
+require "stringio"
+require "uri"
 
 # modules
-require_relative "informers/feature_extraction"
-require_relative "informers/fill_mask"
-require_relative "informers/ner"
-require_relative "informers/question_answering"
-require_relative "informers/sentiment_analysis"
-require_relative "informers/text_generation"
-require_relative "informers/version"
+require_relative "informers/utils/core"
+require_relative "informers/utils/hub"
+require_relative "informers/utils/math"
+require_relative "informers/utils/tensor"
+require_relative "informers/configs"
+require_relative "informers/env"
+require_relative "informers/model"
+require_relative "informers/models"
+require_relative "informers/tokenizers"
+require_relative "informers/pipelines"
+
+module Informers
+  class Error < StandardError; end
+
+  class Todo < Error
+    def message
+      "not implemented yet"
+    end
+  end
+end

lib/informers/configs.rb

@@ -0,0 +1,48 @@
+module Informers
+  class PretrainedConfig
+    attr_reader :model_type, :problem_type, :id2label
+
+    def initialize(config_json)
+      @is_encoder_decoder = false
+
+      @model_type = config_json["model_type"]
+      @problem_type = config_json["problem_type"]
+      @id2label = config_json["id2label"]
+    end
+
+    def [](key)
+      instance_variable_get("@#{key}")
+    end
+
+    def self.from_pretrained(
+      pretrained_model_name_or_path,
+      progress_callback: nil,
+      config: nil,
+      cache_dir: nil,
+      local_files_only: false,
+      revision: "main",
+      **kwargs
+    )
+      data = config || load_config(
+        pretrained_model_name_or_path,
+        progress_callback:,
+        config:,
+        cache_dir:,
+        local_files_only:,
+        revision:
+      )
+      new(data)
+    end
+
+    def self.load_config(pretrained_model_name_or_path, **options)
+      info = Utils::Hub.get_model_json(pretrained_model_name_or_path, "config.json", true, **options)
+      info
+    end
+  end
+
+  class AutoConfig
+    def self.from_pretrained(...)
+      PretrainedConfig.from_pretrained(...)
+    end
+  end
+end

lib/informers/env.rb

@@ -0,0 +1,14 @@
+module Informers
+  CACHE_HOME = ENV.fetch("XDG_CACHE_HOME", File.join(ENV.fetch("HOME"), ".cache"))
+  DEFAULT_CACHE_DIR = File.expand_path(File.join(CACHE_HOME, "informers"))
+
+  class << self
+    attr_accessor :allow_remote_models, :remote_host, :remote_path_template, :cache_dir
+  end
+
+  self.allow_remote_models = ENV["INFORMERS_OFFLINE"].to_s.empty?
+  self.remote_host = "https://huggingface.co/"
+  self.remote_path_template = "{model}/resolve/{revision}/"
+
+  self.cache_dir = DEFAULT_CACHE_DIR
+end