llama_cpp 0.1.4 → 0.2.1

data/ext/llama_cpp/llama_cpp.cpp

@@ -4,6 +4,7 @@
 VALUE rb_mLLaMACpp;
 VALUE rb_cLLaMAContext;
 VALUE rb_cLLaMAContextParams;
+VALUE rb_cLLaMAModelQuantizeParams;
 VALUE rb_cLLaMATokenData;
 VALUE rb_cLLaMATokenDataArray;
 
@@ -292,6 +293,15 @@ public:
     // rb_define_method(rb_cLLaMAContextParams, "initialize", RUBY_METHOD_FUNC(_llama_context_params_init), 0);
     rb_define_method(rb_cLLaMAContextParams, "n_ctx=", RUBY_METHOD_FUNC(_llama_context_params_set_n_ctx), 1);
     rb_define_method(rb_cLLaMAContextParams, "n_ctx", RUBY_METHOD_FUNC(_llama_context_params_get_n_ctx), 0);
+    rb_define_method(rb_cLLaMAContextParams, "n_batch=", RUBY_METHOD_FUNC(_llama_context_params_set_n_batch), 1);
+    rb_define_method(rb_cLLaMAContextParams, "n_batch", RUBY_METHOD_FUNC(_llama_context_params_get_n_batch), 0);
+    rb_define_method(rb_cLLaMAContextParams, "n_gpu_layers=", RUBY_METHOD_FUNC(_llama_context_params_set_n_gpu_layers), 1);
+    rb_define_method(rb_cLLaMAContextParams, "n_gpu_layers", RUBY_METHOD_FUNC(_llama_context_params_get_n_gpu_layers), 0);
+    rb_define_method(rb_cLLaMAContextParams, "main_gpu=", RUBY_METHOD_FUNC(_llama_context_params_set_main_gpu), 1);
+    rb_define_method(rb_cLLaMAContextParams, "main_gpu", RUBY_METHOD_FUNC(_llama_context_params_get_main_gpu), 0);
+    rb_define_method(rb_cLLaMAContextParams, "tensor_split", RUBY_METHOD_FUNC(_llama_context_params_get_tensor_split), 0);
+    rb_define_method(rb_cLLaMAContextParams, "low_vram=", RUBY_METHOD_FUNC(_llama_context_params_set_low_vram), 1);
+    rb_define_method(rb_cLLaMAContextParams, "low_vram", RUBY_METHOD_FUNC(_llama_context_params_get_low_vram), 0);
     rb_define_method(rb_cLLaMAContextParams, "seed=", RUBY_METHOD_FUNC(_llama_context_params_set_seed), 1);
     rb_define_method(rb_cLLaMAContextParams, "seed", RUBY_METHOD_FUNC(_llama_context_params_get_seed), 0);
     rb_define_method(rb_cLLaMAContextParams, "f16_kv=", RUBY_METHOD_FUNC(_llama_context_params_set_f16_kv), 1);
@@ -329,6 +339,67 @@ private:
     return INT2NUM(ptr->params.n_ctx);
   };
 
+  // n_batch
+  static VALUE _llama_context_params_set_n_batch(VALUE self, VALUE n_batch) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    ptr->params.n_batch = NUM2INT(n_batch);
+    return INT2NUM(ptr->params.n_batch);
+  };
+
+  static VALUE _llama_context_params_get_n_batch(VALUE self) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    return INT2NUM(ptr->params.n_batch);
+  };
+
+  // n_gpu_layers
+  static VALUE _llama_context_params_set_n_gpu_layers(VALUE self, VALUE n_gpu_layers) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    ptr->params.n_gpu_layers = NUM2INT(n_gpu_layers);
+    return INT2NUM(ptr->params.n_gpu_layers);
+  };
+
+  static VALUE _llama_context_params_get_n_gpu_layers(VALUE self) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    return INT2NUM(ptr->params.n_gpu_layers);
+  };
+
+  // main_gpu
+  static VALUE _llama_context_params_set_main_gpu(VALUE self, VALUE main_gpu) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    ptr->params.main_gpu = NUM2INT(main_gpu);
+    return INT2NUM(ptr->params.main_gpu);
+  };
+
+  static VALUE _llama_context_params_get_main_gpu(VALUE self) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    return INT2NUM(ptr->params.main_gpu);
+  };
+
+  // tensor_split
+  static VALUE _llama_context_params_get_tensor_split(VALUE self) {
+    if (LLAMA_MAX_DEVICES < 1) {
+      return rb_ary_new();
+    }
+    VALUE ret = rb_ary_new2(LLAMA_MAX_DEVICES);
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    for (size_t i = 0; i < LLAMA_MAX_DEVICES; i++) {
+      rb_ary_store(ret, i, DBL2NUM(ptr->params.tensor_split[i]));
+    }
+    return ret;
+  };
+
+  // low_vram
+  static VALUE _llama_context_params_set_low_vram(VALUE self, VALUE low_vram) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    ptr->params.low_vram = low_vram == Qtrue ? true : false;
+    return ptr->params.low_vram ? Qtrue : Qfalse;
+  };
+
+  static VALUE _llama_context_params_get_low_vram(VALUE self) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    return ptr->params.low_vram ? Qtrue : Qfalse;
+  };
+
   // seed
   static VALUE _llama_context_params_set_seed(VALUE self, VALUE seed) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
@@ -424,6 +495,121 @@ const rb_data_type_t RbLLaMAContextParams::llama_context_params_type = {
   RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+class LLaMAModelQuantizeParamsWrapper {
+public:
+  llama_model_quantize_params params;
+
+  LLaMAModelQuantizeParamsWrapper() : params(llama_model_quantize_default_params()){};
+
+  ~LLaMAModelQuantizeParamsWrapper(){};
+};
+
+class RbLLaMAModelQuantizeParams {
+public:
+  static VALUE llama_model_quantize_params_alloc(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr = (LLaMAModelQuantizeParamsWrapper*)ruby_xmalloc(sizeof(LLaMAModelQuantizeParamsWrapper));
+    new (ptr) LLaMAModelQuantizeParamsWrapper();
+    return TypedData_Wrap_Struct(self, &llama_model_quantize_params_type, ptr);
+  };
+
+  static void llama_model_quantize_params_free(void* ptr) {
+    ((LLaMAModelQuantizeParamsWrapper*)ptr)->~LLaMAModelQuantizeParamsWrapper();
+    ruby_xfree(ptr);
+  };
+
+  static size_t llama_model_quantize_params_size(const void* ptr) {
+    return sizeof(*((LLaMAModelQuantizeParamsWrapper*)ptr));
+  };
+
+  static LLaMAModelQuantizeParamsWrapper* get_llama_model_quantize_params(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr;
+    TypedData_Get_Struct(self, LLaMAModelQuantizeParamsWrapper, &llama_model_quantize_params_type, ptr);
+    return ptr;
+  };
+
+  static void define_class(VALUE outer) {
+    rb_cLLaMAModelQuantizeParams = rb_define_class_under(outer, "ModelQuantizeParams", rb_cObject);
+    rb_define_alloc_func(rb_cLLaMAModelQuantizeParams, llama_model_quantize_params_alloc);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "n_thread=", RUBY_METHOD_FUNC(_llama_model_quantize_params_set_n_thread), 1);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "n_thread", RUBY_METHOD_FUNC(_llama_model_quantize_params_get_n_thread), 0);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "ftype=", RUBY_METHOD_FUNC(_llama_model_quantize_params_set_ftype), 1);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "ftype", RUBY_METHOD_FUNC(_llama_model_quantize_params_get_ftype), 0);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "allow_requantize=", RUBY_METHOD_FUNC(_llama_model_quantize_params_set_allow_requantize), 1);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "allow_requantize", RUBY_METHOD_FUNC(_llama_model_quantize_params_get_allow_requantize), 0);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "quantize_output_tensor=", RUBY_METHOD_FUNC(_llama_model_quantize_params_set_quantize_output_tensor), 1);
+    rb_define_method(rb_cLLaMAModelQuantizeParams, "quantize_output_tensor", RUBY_METHOD_FUNC(_llama_model_quantize_params_get_quantize_output_tensor), 0);
+  };
+
+private:
+  static const rb_data_type_t llama_model_quantize_params_type;
+
+  // n_thread
+  static VALUE _llama_model_quantize_params_set_n_thread(VALUE self, VALUE n_thread) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    ptr->params.nthread = NUM2INT(n_thread);
+    return INT2NUM(ptr->params.nthread);
+  };
+
+  static VALUE _llama_model_quantize_params_get_n_thread(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    return INT2NUM(ptr->params.nthread);
+  };
+
+  // ftype
+  static VALUE _llama_model_quantize_params_set_ftype(VALUE self, VALUE ftype) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    ptr->params.ftype = static_cast<enum llama_ftype>(NUM2INT(ftype));
+    return INT2NUM(ptr->params.ftype);
+  };
+
+  static VALUE _llama_model_quantize_params_get_ftype(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    return INT2NUM(ptr->params.ftype);
+  };
+
+  // allow_requantize
+  static VALUE _llama_model_quantize_params_set_allow_requantize(VALUE self, VALUE allow_requantize) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    if (NIL_P(allow_requantize) || allow_requantize == Qfalse) {
+      ptr->params.allow_requantize = false;
+    } else {
+      ptr->params.allow_requantize = true;
+    }
+    return ptr->params.allow_requantize ? Qtrue : Qfalse;
+  };
+
+  static VALUE _llama_model_quantize_params_get_allow_requantize(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    return ptr->params.allow_requantize ? Qtrue : Qfalse;
+  };
+
+  // quantize_output_tensor
+  static VALUE _llama_model_quantize_params_set_quantize_output_tensor(VALUE self, VALUE quantize_output_tensor) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    if (NIL_P(quantize_output_tensor) || quantize_output_tensor == Qfalse) {
+      ptr->params.quantize_output_tensor = false;
+    } else {
+      ptr->params.quantize_output_tensor = true;
+    }
+    return ptr->params.quantize_output_tensor ? Qtrue : Qfalse;
+  };
+
+  static VALUE _llama_model_quantize_params_get_quantize_output_tensor(VALUE self) {
+    LLaMAModelQuantizeParamsWrapper* ptr = get_llama_model_quantize_params(self);
+    return ptr->params.quantize_output_tensor ? Qtrue : Qfalse;
+  };
+};
+
+const rb_data_type_t RbLLaMAModelQuantizeParams::llama_model_quantize_params_type = {
+  "RbLLaMAModelQuantizeParams",
+  { NULL,
+    RbLLaMAModelQuantizeParams::llama_model_quantize_params_free,
+    RbLLaMAModelQuantizeParams::llama_model_quantize_params_size },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+};
+
 class LLaMAContextWrapper {
 public:
   struct llama_context* ctx;
@@ -465,9 +651,11 @@ public:
     rb_define_alloc_func(rb_cLLaMAContext, llama_context_alloc);
     rb_define_method(rb_cLLaMAContext, "initialize", RUBY_METHOD_FUNC(_llama_context_initialize), -1);
     rb_define_method(rb_cLLaMAContext, "eval", RUBY_METHOD_FUNC(_llama_context_eval), -1);
+    rb_define_method(rb_cLLaMAContext, "eval_export", RUBY_METHOD_FUNC(_llama_context_eval_export), 1);
     rb_define_method(rb_cLLaMAContext, "tokenize", RUBY_METHOD_FUNC(_llama_context_tokenize), -1);
     rb_define_method(rb_cLLaMAContext, "logits", RUBY_METHOD_FUNC(_llama_context_logits), 0);
     rb_define_method(rb_cLLaMAContext, "embeddings", RUBY_METHOD_FUNC(_llama_context_embeddings), 0);
+    rb_define_method(rb_cLLaMAContext, "vocab", RUBY_METHOD_FUNC(_llama_context_vocab), -1);
     rb_define_method(rb_cLLaMAContext, "token_to_str", RUBY_METHOD_FUNC(_llama_context_token_to_str), 1);
     rb_define_method(rb_cLLaMAContext, "n_vocab", RUBY_METHOD_FUNC(_llama_context_n_vocab), 0);
     rb_define_method(rb_cLLaMAContext, "n_ctx", RUBY_METHOD_FUNC(_llama_context_n_ctx), 0);
@@ -517,7 +705,7 @@ private:
       return Qnil;
     }
     if (!rb_obj_is_kind_of(kw_values[1], rb_cLLaMAContextParams)) {
-      rb_raise(rb_eArgError, "params must be a LLaMAContextParams");
+      rb_raise(rb_eArgError, "params must be a ContextParams");
       return Qnil;
     }
 
@@ -599,6 +787,24 @@ private:
     return Qnil;
   };
 
+  static VALUE _llama_context_eval_export(VALUE self, VALUE fname_) {
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    if (!RB_TYPE_P(fname_, T_STRING)) {
+      rb_raise(rb_eArgError, "fname must be a string");
+      return Qnil;
+    }
+    const char* fname = StringValueCStr(fname_);
+    if (llama_eval_export(ptr->ctx, fname) != 0) {
+      return Qfalse;
+    }
+    RB_GC_GUARD(fname_);
+    return Qtrue;
+  };
+
   static VALUE _llama_context_tokenize(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
     ID kw_table[3] = { rb_intern("text"), rb_intern("n_max_tokens"), rb_intern("add_bos") };
@@ -705,6 +911,43 @@ private:
     return output;
   };
 
+  static VALUE _llama_context_vocab(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("capacity") };
+    VALUE kw_values[1] = { Qundef };
+    rb_scan_args(argc, argv, ":", &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "capacity must be an integer");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+
+    const int capacity = NUM2INT(kw_values[0]);
+    std::vector<const char*> strings;
+    std::vector<float> scores;
+    int n_vocab = llama_n_vocab(ptr->ctx);
+    strings.resize(n_vocab, NULL);
+    scores.resize(n_vocab, 0);
+
+    n_vocab = llama_get_vocab(ptr->ctx, strings.data(), scores.data(), capacity);
+
+    VALUE ret_strings = rb_ary_new();
+    VALUE ret_scores = rb_ary_new();
+    for (int i = 0; i < n_vocab; i++) {
+      rb_ary_push(ret_strings, rb_utf8_str_new_cstr(strings[i]));
+      rb_ary_push(ret_scores, DBL2NUM(static_cast<double>(scores[i])));
+    }
+
+    return rb_ary_new_from_args(2, ret_strings, ret_scores);
+  };
+
   static VALUE _llama_context_n_vocab(VALUE self) {
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
@@ -1428,10 +1671,10 @@ static VALUE rb_llama_llama_init_backend(VALUE self) {
 
 static VALUE rb_llama_model_quantize(int argc, VALUE* argv, VALUE self) {
   VALUE kw_args = Qnil;
-  ID kw_table[4] = { rb_intern("input_path"), rb_intern("output_path"), rb_intern("ftype"), rb_intern("n_threads") };
-  VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
+  ID kw_table[3] = { rb_intern("input_path"), rb_intern("output_path"), rb_intern("params") };
+  VALUE kw_values[3] = { Qundef, Qundef, Qundef };
   rb_scan_args(argc, argv, ":", &kw_args);
-  rb_get_kwargs(kw_args, kw_table, 3, 1, kw_values);
+  rb_get_kwargs(kw_args, kw_table, 3, 0, kw_values);
 
   if (!RB_TYPE_P(kw_values[0], T_STRING)) {
     rb_raise(rb_eArgError, "input_path must be a string");
@@ -1441,21 +1684,16 @@ static VALUE rb_llama_model_quantize(int argc, VALUE* argv, VALUE self) {
     rb_raise(rb_eArgError, "output_path must be a string");
     return Qnil;
   }
-  if (!RB_INTEGER_TYPE_P(kw_values[2])) {
-    rb_raise(rb_eArgError, "ftype must be an integer");
-    return Qnil;
-  }
-  if (kw_values[3] != Qundef && !RB_INTEGER_TYPE_P(kw_values[3])) {
-    rb_raise(rb_eArgError, "n_threads must be an integer");
+  if (!rb_obj_is_kind_of(kw_values[2], rb_cLLaMAModelQuantizeParams)) {
+    rb_raise(rb_eArgError, "params must be a ModelQuantizeParams");
     return Qnil;
   }
 
   const char* input_path = StringValueCStr(kw_values[0]);
   const char* output_path = StringValueCStr(kw_values[1]);
-  const int ftype = NUM2INT(kw_values[2]);
-  const int n_threads = kw_values[3] == Qundef ? 1 : NUM2INT(kw_values[3]);
+  LLaMAModelQuantizeParamsWrapper* wrapper = RbLLaMAModelQuantizeParams::get_llama_model_quantize_params(kw_values[2]);
 
-  if (llama_model_quantize(input_path, output_path, (llama_ftype)ftype, n_threads) != 0) {
+  if (llama_model_quantize(input_path, output_path, &(wrapper->params)) != 0) {
     rb_raise(rb_eRuntimeError, "Failed to quantize model");
     return Qnil;
   }
@@ -1505,6 +1743,8 @@ extern "C" void Init_llama_cpp(void) {
   rb_define_module_function(rb_mLLaMACpp, "mmap_supported?", rb_llama_mmap_supported, 0);
   rb_define_module_function(rb_mLLaMACpp, "mlock_supported?", rb_llama_mlock_supported, 0);
 
+  rb_define_const(rb_mLLaMACpp, "LLAMA_MAX_DEVICES", INT2NUM(LLAMA_MAX_DEVICES));
+
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_ALL_F32", INT2NUM(LLAMA_FTYPE_ALL_F32));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_F16", INT2NUM(LLAMA_FTYPE_MOSTLY_F16));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_0", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_0));
@@ -1513,6 +1753,15 @@ extern "C" void Init_llama_cpp(void) {
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q8_0", INT2NUM(LLAMA_FTYPE_MOSTLY_Q8_0));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q5_0", INT2NUM(LLAMA_FTYPE_MOSTLY_Q5_0));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q5_1", INT2NUM(LLAMA_FTYPE_MOSTLY_Q5_1));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q2_K", INT2NUM(LLAMA_FTYPE_MOSTLY_Q2_K));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q3_K_S", INT2NUM(LLAMA_FTYPE_MOSTLY_Q3_K_S));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q3_K_M", INT2NUM(LLAMA_FTYPE_MOSTLY_Q3_K_M));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q3_K_L", INT2NUM(LLAMA_FTYPE_MOSTLY_Q3_K_L));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_K_S", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_K_S));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_K_M", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_K_M));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q5_K_S", INT2NUM(LLAMA_FTYPE_MOSTLY_Q5_K_S));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q5_K_M", INT2NUM(LLAMA_FTYPE_MOSTLY_Q5_K_M));
+  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q6_K", INT2NUM(LLAMA_FTYPE_MOSTLY_Q6_K));
 
   std::stringstream ss_magic;
   ss_magic << std::showbase << std::hex << LLAMA_FILE_MAGIC_GGJT;