gpt_neox_client 0.1.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d77a992f9cbba2e7a91141c859194cc0a200b9c5dd4e58aeedc51482ec75b8c0
-  data.tar.gz: 302c37b125a0875463859b62fb7918b47c38cd521ad444fe758a36ba06e27ade
+  metadata.gz: f7d90a7d8178a4974871638030ee67311c7ec38c169810c0c4b583d4cd1d697a
+  data.tar.gz: 3c85344089c5f1048524b1163cf956c4b723a0c29b771706a3b514c7bc2088aa
 SHA512:
-  metadata.gz: a3ab1eb43db87f08e24fb16181d83feeb4b1421fd601165a233cc48eccfd0403de2e1d98042e89e259c843e5c3021e3931d57f4742ee20ae8c7e55cf45f6c0d4
-  data.tar.gz: 88f961e5a901ea5896486b4612ab3ecc8ebcbad12726fd76700bb1248e31c11716ae4a7248592f7657fb056e36af19f466f95860862a4d165fc026df3d4cb04f
+  metadata.gz: 7f0ac814530db33cd077505b093a5ec1fa1cfe715541c5210d81c47ce34e927128422c503a5f6f55474d5200c59e566e7d41c10243518897a98e4535d0588f5d
+  data.tar.gz: efad11b4aebd6b07070ab6d9b043f51832a24c23a8a6adde52167df8e0128339f063d1375564ded5f775e953ef3543b1f1eda20dbf3ab0313a68380396298b6a

data/CHANGELOG.md

@@ -1,5 +1,20 @@
 ## [Unreleased]
 
-## [0.1.0] - 2023-09-xx
+## [0.3.0] - 2023-09-06
+
+- Add `embeddings` method.
+  ```ruby
+  require 'gpt_neox_client'
+
+  client = GPTNeoXClient.new(path: '/path/to/ggml-model-f16.bin', , seed: 123456789, n_threads: 8)
+  embd = client.embeddings('Hello, world.', normalize: true)
+  ```
+
+## [0.2.0] - 2023-09-02
+
+- Add Accelerate framework and Metal build option for macOS.
+- Add OpenBLAS build option for platforms other than macOS.
+
+## [0.1.0] - 2023-09-01
 
 - Initial release

data/README.md

@@ -38,15 +38,15 @@ japanese-large-lm-3.6b-instruction-sft/ggml-model-f16.bin
 ```
 
 ```ruby
-require "gpt_neox_client"
+require 'gpt_neox_client'
 
-client = GPTNeoXClient.new('japanese-large-lm-3.6b-instruction-sft/ggml-model-f16.bin', seed: 123456789, n_threads: 4)
+client = GPTNeoXClient.new(path: 'japanese-large-lm-3.6b-instruction-sft/ggml-model-f16.bin', seed: 123456789, n_threads: 4)
 puts client.completions(
   'ユーザー:四国の県名を全て列挙してください。<0x0A>システム:',
   top_p: 0.9,
   top_k: 1,
   temperature: 0.7
-).gsub("<0x0A>", "\n").gsub("</s>", " ")
+).gsub('<0x0A>', "\n").gsub('</s>', '')
 #
 # ユーザー:四国の県名を全て列挙してください。
 # システム:徳島県、香川県、愛媛県、高知県

data/ext/gpt_neox_client/extconf.rb

@@ -22,4 +22,30 @@ $INCFLAGS << ' -I$(srcdir)/src/ggml'
 $VPATH << '$(srcdir)/src'
 $VPATH << '$(srcdir)/src/ggml'
 
+if RUBY_PLATFORM.match?(/darwin|linux|bsd/) && try_compile('#include <stdio.h>', '-pthread')
+  $CFLAGS << ' -pthread'
+  $CXXFLAGS << ' -pthread'
+end
+
+if RUBY_PLATFORM.match?(/darwin/)
+  if have_framework('Accelerate')
+    $CFLAGS << ' -DGGML_USE_ACCELERATE'
+  else
+    warning 'Accelerate framework is not found.'
+  end
+end
+
+$CFLAGS << ' -DGGML_USE_OPENBLAS' if !RUBY_PLATFORM.match?(/darwin/) && (have_library('openblas') && have_header('cblas.h'))
+
 create_makefile('gpt_neox_client/gpt_neox_client')
+
+if RUBY_PLATFORM.match?(/darwin/)
+  File.open('Makefile', 'a') do |f|
+    f.puts "\nggml-metal.o: ggml-metal.m ggml-metal.h"
+    f.puts "\t$(CC) $(CFLAGS) -c $< -o $@"
+  end
+
+  metal_path = File.expand_path("#{__dir__}/src/ggml/ggml-metal.metal")
+  dest_path = File.expand_path("#{__dir__}/../../lib/gpt_neox_client/")
+  FileUtils.cp(metal_path, dest_path)
+end

data/ext/gpt_neox_client/gpt_neox_client.cpp

@@ -196,9 +196,10 @@ static VALUE gpt_neox_client_completions(int argc, VALUE* argv, VALUE self) {
   const int n_predict = std::min(n_predict_, model->hparams.n_ctx - static_cast<int>(embd_inp.size()));
 
   const int n_threads = NUM2INT(rb_iv_get(self, "@n_threads"));
+  std::vector<float> embedding;
   std::vector<float> logits;
   size_t mem_per_token = 0;
-  gpt_neox_eval(*model, n_threads, 0, { 0, 1, 2, 3 }, logits, mem_per_token);
+  gpt_neox_eval(*model, n_threads, 0, { 0, 1, 2, 3 }, embedding, logits, mem_per_token);
 
   int n_past = 0;
   int n_consumed = 0;
@@ -208,10 +209,11 @@ static VALUE gpt_neox_client_completions(int argc, VALUE* argv, VALUE self) {
   std::mt19937 rng(seed);
   std::vector<gpt_vocab::id> embd;
   std::vector<int32_t> last_n_tokens(model->hparams.n_ctx, 0);
+  gpt_vocab::id token_eos = vocab->token_to_id["</s>"];
 
   while (n_sampled < n_predict) {
     if (embd.size() > 0) {
-      if (!gpt_neox_eval(*model, n_threads, n_past, embd, logits, mem_per_token)) {
+      if (!gpt_neox_eval(*model, n_threads, n_past, embd, embedding, logits, mem_per_token)) {
         rb_raise(rb_eRuntimeError, "failed to predict.");
         return Qnil;
       }
@@ -240,13 +242,73 @@ static VALUE gpt_neox_client_completions(int argc, VALUE* argv, VALUE self) {
     }
 
     for (auto id : embd) completions += vocab->id_to_token[id];
-    if (embd.back() == 0) break;
+    if (!embd.empty() && embd.back() == token_eos) break;
   }
 
   RB_GC_GUARD(prompt_);
   return rb_utf8_str_new_cstr(completions.c_str());
 }
 
+static VALUE gpt_neox_client_embeddings(int argc, VALUE* argv, VALUE self) {
+  VALUE prompt_ = Qnil;
+  VALUE kw_args = Qnil;
+  rb_scan_args(argc, argv, "1:", &prompt_, &kw_args);
+
+  ID kw_table[2] = { rb_intern("n_batch"), rb_intern("normalize") };
+  VALUE kw_values[2] = { Qundef, Qundef };
+  rb_get_kwargs(kw_args, kw_table, 0, 2, kw_values);
+
+  if (kw_values[0] != Qundef && !RB_INTEGER_TYPE_P(kw_values[0])) {
+    rb_raise(rb_eArgError, "n_batch must be an integer");
+    return Qnil;
+  }
+
+  std::string prompt(StringValueCStr(prompt_));
+  const int n_batch = kw_values[0] != Qundef ? NUM2INT(kw_values[0]) : 8;
+  const bool normalize = kw_values[1] != Qundef ? RTEST(kw_values[1]) : false;
+
+  gpt_neox_model* model = RbGPTNeoXModel::get_gpt_neox_model(rb_iv_get(self, "@model"));
+  gpt_vocab* vocab = RbGPTVocab::get_gpt_vocab(rb_iv_get(self, "@vocab"));
+  const int n_threads = NUM2INT(rb_iv_get(self, "@n_threads"));
+
+  std::vector<gpt_vocab::id> embd_inp = gpt_tokenize(*vocab, prompt);
+
+  if (embd_inp.size() > model->hparams.n_ctx) {
+    rb_raise(rb_eArgError, "prompt is too long");
+    return Qnil;
+  }
+
+  std::vector<float> embedding;
+  std::vector<float> logits;
+  size_t mem_per_token = 0;
+  gpt_neox_eval(*model, n_threads, 0, { 0, 1, 2, 3 }, embedding, logits, mem_per_token);
+
+  int n_past = 0;
+  std::vector<gpt_vocab::id> embd;
+  while (!embd_inp.empty()) {
+    const int n_tokens = std::min(n_batch, static_cast<int>(embd_inp.size()));
+    embd.insert(embd.end(), embd_inp.begin(), embd_inp.begin() + n_tokens);
+    if (!gpt_neox_eval(*model, n_threads, n_past, embd, embedding, logits, mem_per_token)) {
+      rb_raise(rb_eRuntimeError, "failed to predict.");
+      return Qnil;
+    }
+    n_past += n_tokens;
+    embd.clear();
+    embd_inp.erase(embd_inp.begin(), embd_inp.begin() + n_tokens);
+  }
+
+  if (normalize) {
+    const float norm = std::sqrt(std::inner_product(embedding.begin(), embedding.end(), embedding.begin(), 0.0f));
+    for (auto& v : embedding) v /= norm;
+  }
+
+  VALUE res = rb_ary_new2(embedding.size());
+  for (size_t i = 0; i < embedding.size(); i++) rb_ary_store(res, i, DBL2NUM(embedding[i]));
+
+  RB_GC_GUARD(prompt_);
+  return res;
+}
+
 extern "C" void Init_gpt_neox_client(void) {
   /**
    * Document-class: GPTNeoXClient
@@ -289,6 +351,22 @@ extern "C" void Init_gpt_neox_client(void) {
    * @return [String]
    */
   rb_define_method(rb_cGPTNeoXClient, "completions", RUBY_METHOD_FUNC(gpt_neox_client_completions), -1);
+  /**
+   * Generates embeddings.
+   *
+   * @example
+   *   require "gpt_neox_client"
+   *
+   *   client = GPTNeoXClient.new("gpt-neox-f16.bin")
+   *   client.embeddings("Hello, my name is")
+   *
+   * @overload embeddings(text, n_batch: 8, normalize: false)
+   *   @param [String] text The text.
+   *   @param [Integer] n_batch The number of tokens to evalauate at once.
+   *   @param [Boolean] normalize The flag to normalize the embeddings.
+   * @return [Array<Float>]
+   */
+  rb_define_method(rb_cGPTNeoXClient, "embeddings", RUBY_METHOD_FUNC(gpt_neox_client_embeddings), -1);
   /**
    * Returns the path to the model.
    * @return [String]

data/ext/gpt_neox_client/src/main.cpp

@@ -433,6 +433,7 @@ bool gpt_neox_eval(
         const int n_threads,
         const int n_past,
         const std::vector<gpt_vocab::id> & embd_inp,
+              std::vector<float>         & embd_d,
               std::vector<float>         & embd_w,
               size_t                     & mem_per_token) {
     const int N = embd_inp.size();
@@ -657,6 +658,10 @@ bool gpt_neox_eval(
     //embd_w.resize(n_vocab*N);
     //memcpy(embd_w.data(), ggml_get_data(inpL), sizeof(float)*n_vocab*N);
 
+    embd_d.resize(n_embd);
+    struct ggml_tensor* embeddings = gf.nodes[gf.n_nodes - 2];
+    memcpy(embd_d.data(), (float*)ggml_get_data(embeddings) + (n_embd * (N - 1)), sizeof(float)*n_embd);
+
     // return result for just the last token
     embd_w.resize(n_vocab);
     memcpy(embd_w.data(), (float *) ggml_get_data(inpL) + (n_vocab*(N-1)), sizeof(float)*n_vocab);

data/lib/gpt_neox_client/version.rb

@@ -3,5 +3,5 @@
 # GPTNeoXClient is a Ruby client for GPT-NeoX.
 class GPTNeoXClient
   # The version of GPTNeoXClient you are using.
-  VERSION = '0.1.0'
+  VERSION = '0.3.0'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: gpt_neox_client
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.3.0
 platform: ruby
 authors:
 - yoshoku
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-09-01 00:00:00.000000000 Z
+date: 2023-09-06 00:00:00.000000000 Z
 dependencies: []
 description: gpt_neox_client is a simple client for GPT-NeoX.
 email: