llama_cpp 0.0.7 → 0.1.0

data/ext/llama_cpp/llama_cpp.cpp

@@ -4,6 +4,255 @@
 VALUE rb_mLLaMACpp;
 VALUE rb_cLLaMAContext;
 VALUE rb_cLLaMAContextParams;
+VALUE rb_cLLaMATokenData;
+VALUE rb_cLLaMATokenDataArray;
+
+class LLaMATokenDataWrapper {
+public:
+  llama_token_data data;
+
+  LLaMATokenDataWrapper() {
+    data.id = 0;
+    data.logit = 0.0;
+    data.p = 0.0;
+  };
+
+  ~LLaMATokenDataWrapper(){};
+};
+
+class RbLLaMATokenData {
+public:
+  static VALUE llama_token_data_alloc(VALUE self) {
+    LLaMATokenDataWrapper* ptr = (LLaMATokenDataWrapper*)ruby_xmalloc(sizeof(LLaMATokenDataWrapper));
+    new (ptr) LLaMATokenDataWrapper();
+    return TypedData_Wrap_Struct(self, &llama_token_data_type, ptr);
+  };
+
+  static void llama_token_data_free(void* ptr) {
+    ((LLaMATokenDataWrapper*)ptr)->~LLaMATokenDataWrapper();
+    ruby_xfree(ptr);
+  };
+
+  static size_t llama_token_data_size(const void* ptr) {
+    return sizeof(*((LLaMATokenDataWrapper*)ptr));
+  };
+
+  static LLaMATokenDataWrapper* get_llama_token_data(VALUE self) {
+    LLaMATokenDataWrapper* ptr;
+    TypedData_Get_Struct(self, LLaMATokenDataWrapper, &llama_token_data_type, ptr);
+    return ptr;
+  };
+
+  static void define_class(VALUE outer) {
+    rb_cLLaMATokenData = rb_define_class_under(outer, "TokenData", rb_cObject);
+    rb_define_alloc_func(rb_cLLaMATokenData, llama_token_data_alloc);
+    rb_define_method(rb_cLLaMATokenData, "initialize", RUBY_METHOD_FUNC(_llama_token_data_init), -1);
+    rb_define_method(rb_cLLaMATokenData, "id=", RUBY_METHOD_FUNC(_llama_token_data_set_id), 1);
+    rb_define_method(rb_cLLaMATokenData, "id", RUBY_METHOD_FUNC(_llama_token_data_get_id), 0);
+    rb_define_method(rb_cLLaMATokenData, "logit=", RUBY_METHOD_FUNC(_llama_token_data_set_logit), 1);
+    rb_define_method(rb_cLLaMATokenData, "logit", RUBY_METHOD_FUNC(_llama_token_data_get_logit), 0);
+    rb_define_method(rb_cLLaMATokenData, "p=", RUBY_METHOD_FUNC(_llama_token_data_set_p), 1);
+    rb_define_method(rb_cLLaMATokenData, "p", RUBY_METHOD_FUNC(_llama_token_data_get_p), 0);
+  }
+
+private:
+  static const rb_data_type_t llama_token_data_type;
+
+  static VALUE _llama_token_data_init(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[3] = { rb_intern("id"), rb_intern("logit"), rb_intern("p") };
+    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
+    rb_scan_args(argc, argv, ":", &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 3, 0, kw_values);
+
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "id must be an integer");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "logit must be a float");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[2])) {
+      rb_raise(rb_eArgError, "p must be a float");
+      return Qnil;
+    }
+
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    new (ptr) LLaMATokenDataWrapper();
+
+    ptr->data.id = NUM2INT(kw_values[0]);
+    ptr->data.logit = NUM2DBL(kw_values[1]);
+    ptr->data.p = NUM2DBL(kw_values[2]);
+
+    return self;
+  }
+
+  // id
+  static VALUE _llama_token_data_set_id(VALUE self, VALUE id) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    ptr->data.id = NUM2INT(id);
+    return INT2NUM(ptr->data.id);
+  };
+
+  static VALUE _llama_token_data_get_id(VALUE self) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    return INT2NUM(ptr->data.id);
+  };
+
+  // logit
+  static VALUE _llama_token_data_set_logit(VALUE self, VALUE logit) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    ptr->data.logit = NUM2DBL(logit);
+    return DBL2NUM(ptr->data.logit);
+  };
+
+  static VALUE _llama_token_data_get_logit(VALUE self) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    return DBL2NUM(ptr->data.logit);
+  };
+
+  // p
+  static VALUE _llama_token_data_set_p(VALUE self, VALUE p) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    ptr->data.p = NUM2DBL(p);
+    return DBL2NUM(ptr->data.p);
+  };
+
+  static VALUE _llama_token_data_get_p(VALUE self) {
+    LLaMATokenDataWrapper* ptr = get_llama_token_data(self);
+    return DBL2NUM(ptr->data.p);
+  };
+};
+
+const rb_data_type_t RbLLaMATokenData::llama_token_data_type = {
+  "RbLLaMATokenData",
+  { NULL,
+    RbLLaMATokenData::llama_token_data_free,
+    RbLLaMATokenData::llama_token_data_size },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+class LLaMATokenDataArrayWrapper {
+public:
+  llama_token_data_array array;
+
+  LLaMATokenDataArrayWrapper() {
+    array.data = nullptr;
+    array.size = 0;
+    array.sorted = false;
+  };
+
+  ~LLaMATokenDataArrayWrapper() {
+    if (array.data) {
+      ruby_xfree(array.data);
+      array.data = nullptr;
+    }
+  };
+};
+
+class RbLLaMATokenDataArray {
+public:
+  static VALUE llama_token_data_array_alloc(VALUE self) {
+    LLaMATokenDataArrayWrapper* ptr = (LLaMATokenDataArrayWrapper*)ruby_xmalloc(sizeof(LLaMATokenDataArrayWrapper));
+    new (ptr) LLaMATokenDataArrayWrapper();
+    return TypedData_Wrap_Struct(self, &llama_token_data_array_type, ptr);
+  };
+
+  static void llama_token_data_array_free(void* ptr) {
+    ((LLaMATokenDataArrayWrapper*)ptr)->~LLaMATokenDataArrayWrapper();
+    ruby_xfree(ptr);
+  };
+
+  static size_t llama_token_data_array_size(const void* ptr) {
+    return sizeof(*((LLaMATokenDataArrayWrapper*)ptr));
+  };
+
+  static LLaMATokenDataArrayWrapper* get_llama_token_data_array(VALUE self) {
+    LLaMATokenDataArrayWrapper* ptr;
+    TypedData_Get_Struct(self, LLaMATokenDataArrayWrapper, &llama_token_data_array_type, ptr);
+    return ptr;
+  };
+
+  static void define_class(VALUE outer) {
+    rb_cLLaMATokenDataArray = rb_define_class_under(outer, "TokenDataArray", rb_cObject);
+    rb_define_alloc_func(rb_cLLaMATokenDataArray, llama_token_data_array_alloc);
+    rb_define_method(rb_cLLaMATokenDataArray, "initialize", RUBY_METHOD_FUNC(_llama_token_data_array_init), -1);
+    rb_define_method(rb_cLLaMATokenDataArray, "size", RUBY_METHOD_FUNC(_llama_token_data_array_get_size), 0);
+    rb_define_method(rb_cLLaMATokenDataArray, "sorted", RUBY_METHOD_FUNC(_llama_token_data_array_get_sorted), 0);
+  };
+
+private:
+  static const rb_data_type_t llama_token_data_array_type;
+
+  static VALUE _llama_token_data_array_init(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("sorted") };
+    VALUE kw_values[1] = { Qundef };
+    VALUE arr = Qnil;
+    rb_scan_args(argc, argv, "1:", &arr, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 0, 1, kw_values);
+
+    if (!RB_TYPE_P(arr, T_ARRAY)) {
+      rb_raise(rb_eArgError, "1st argument must be an array");
+      return Qnil;
+    }
+    size_t sz_array = RARRAY_LEN(arr);
+    if (sz_array == 0) {
+      rb_raise(rb_eArgError, "array must not be empty");
+      return Qnil;
+    }
+    if (kw_values[0] != Qundef && !RB_TYPE_P(kw_values[0], T_TRUE) && !RB_TYPE_P(kw_values[0], T_FALSE)) {
+      rb_raise(rb_eArgError, "sorted must be a boolean");
+      return Qnil;
+    }
+
+    LLaMATokenDataArrayWrapper* ptr = get_llama_token_data_array(self);
+    new (ptr) LLaMATokenDataArrayWrapper();
+
+    ptr->array.data = (llama_token_data*)ruby_xmalloc(sizeof(llama_token_data) * sz_array);
+    for (size_t i = 0; i < sz_array; ++i) {
+      VALUE el = rb_ary_entry(arr, i);
+      if (!rb_obj_is_kind_of(el, rb_cLLaMATokenData)) {
+        rb_raise(rb_eArgError, "array element must be a TokenData");
+        xfree(ptr->array.data);
+        ptr->array.data = nullptr;
+        return Qnil;
+      }
+      llama_token_data token_data = RbLLaMATokenData::get_llama_token_data(el)->data;
+      ptr->array.data[i].id = token_data.id;
+      ptr->array.data[i].logit = token_data.logit;
+      ptr->array.data[i].p = token_data.p;
+    }
+
+    ptr->array.size = sz_array;
+    ptr->array.sorted = kw_values[0] == Qtrue;
+
+    return self;
+  };
+
+  static VALUE _llama_token_data_array_get_size(VALUE self) {
+    LLaMATokenDataArrayWrapper* ptr = get_llama_token_data_array(self);
+    return SIZET2NUM(ptr->array.size);
+  };
+
+  static VALUE _llama_token_data_array_get_sorted(VALUE self) {
+    LLaMATokenDataArrayWrapper* ptr = get_llama_token_data_array(self);
+    return ptr->array.sorted ? Qtrue : Qfalse;
+  };
+};
+
+const rb_data_type_t RbLLaMATokenDataArray::llama_token_data_array_type = {
+  "RbLLaMATokenDataArray",
+  { NULL,
+    RbLLaMATokenDataArray::llama_token_data_array_free,
+    RbLLaMATokenDataArray::llama_token_data_array_size },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+};
 
 class LLaMAContextParamsWrapper {
 public:
@@ -234,7 +483,6 @@ public:
     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, "token_to_str", RUBY_METHOD_FUNC(_llama_context_token_to_str), 1);
-    rb_define_method(rb_cLLaMAContext, "sample_top_p_top_k", RUBY_METHOD_FUNC(_llama_context_sample_top_p_top_k), -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);
     rb_define_method(rb_cLLaMAContext, "n_embd", RUBY_METHOD_FUNC(_llama_context_n_embd), 0);
@@ -244,6 +492,20 @@ public:
     rb_define_method(rb_cLLaMAContext, "free", RUBY_METHOD_FUNC(_llama_context_free), 0);
     rb_define_method(rb_cLLaMAContext, "load", RUBY_METHOD_FUNC(_llama_context_load), -1);
     rb_define_method(rb_cLLaMAContext, "apply_lora_from_file", RUBY_METHOD_FUNC(_llama_context_apply_lora_from_file), -1);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_token_count", RUBY_METHOD_FUNC(_llama_context_kv_cache_token_count), 0);
+    rb_define_method(rb_cLLaMAContext, "set_rng_seed", RUBY_METHOD_FUNC(_llama_context_set_rng_seed), 1);
+    rb_define_method(rb_cLLaMAContext, "sample_repetition_penalty", RUBY_METHOD_FUNC(_llama_context_sample_repetition_penalty), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_frequency_and_presence_penalties", RUBY_METHOD_FUNC(_llama_context_sample_frequency_and_presence_penalties), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_softmax", RUBY_METHOD_FUNC(_llama_context_sample_softmax), 1);
+    rb_define_method(rb_cLLaMAContext, "sample_top_k", RUBY_METHOD_FUNC(_llama_context_sample_top_k), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_top_p", RUBY_METHOD_FUNC(_llama_context_sample_top_p), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_tail_free", RUBY_METHOD_FUNC(_llama_context_sample_tail_free), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_typical", RUBY_METHOD_FUNC(_llama_context_sample_typical), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_temperature", RUBY_METHOD_FUNC(_llama_context_sample_temperature), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_token_mirostat", RUBY_METHOD_FUNC(_llama_context_sample_token_mirostat), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_token_mirostat_v2", RUBY_METHOD_FUNC(_llama_context_sample_token_mirostat_v2), -1);
+    rb_define_method(rb_cLLaMAContext, "sample_token_greedy", RUBY_METHOD_FUNC(_llama_context_sample_token_greedy), 1);
+    rb_define_method(rb_cLLaMAContext, "sample_token", RUBY_METHOD_FUNC(_llama_context_sample_token), 1);
   };
 
 private:
@@ -448,40 +710,6 @@ private:
     return output;
   };
 
-  static VALUE _llama_context_sample_top_p_top_k(int argc, VALUE* argv, VALUE self) {
-    VALUE last_n_tokens = Qnil;
-    VALUE kw_args = Qnil;
-    ID kw_table[4] = { rb_intern("top_k"), rb_intern("top_p"), rb_intern("temp"), rb_intern("penalty") };
-    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
-    rb_scan_args(argc, argv, "1:", &last_n_tokens, &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 4, 0, kw_values);
-
-    if (!RB_TYPE_P(last_n_tokens, T_ARRAY)) {
-      rb_raise(rb_eArgError, "last_n_tokens must be an Array");
-      return Qnil;
-    }
-
-    const int last_n_tokens_size = RARRAY_LEN(last_n_tokens);
-    const int top_k = NUM2INT(kw_values[0]);
-    const double top_p = NUM2DBL(kw_values[1]);
-    const double temp = NUM2DBL(kw_values[2]);
-    const double penalty = NUM2DBL(kw_values[3]);
-
-    std::vector<llama_token> last_n_tokens_data(last_n_tokens_size);
-    for (int i = 0; i < last_n_tokens_size; i++) {
-      last_n_tokens_data[i] = NUM2INT(rb_ary_entry(last_n_tokens, i));
-    }
-
-    LLaMAContextWrapper* ptr = get_llama_context(self);
-    if (ptr->ctx == NULL) {
-      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
-      return Qnil;
-    }
-    llama_token token = llama_sample_top_p_top_k(ptr->ctx, last_n_tokens_data.data(), last_n_tokens_size, top_k, top_p, temp, penalty);
-
-    return INT2NUM(token);
-  }
-
   static VALUE _llama_context_n_vocab(VALUE self) {
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
@@ -621,6 +849,471 @@ private:
     }
     return Qnil;
   };
+
+  static VALUE _llama_context_kv_cache_token_count(VALUE self) {
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    return INT2NUM(llama_get_kv_cache_token_count(ptr->ctx));
+  };
+
+  static VALUE _llama_context_set_rng_seed(VALUE self, VALUE seed_) {
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    const int seed = NUM2INT(seed_);
+    llama_set_rng_seed(ptr->ctx, seed);
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_repetition_penalty(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("penalty") };
+    VALUE kw_values[1] = { Qundef };
+    VALUE candidates = Qnil;
+    VALUE last_n_tokens = Qnil;
+    rb_scan_args(argc, argv, "2:", &candidates, &last_n_tokens, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "candidates must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_TYPE_P(last_n_tokens, T_ARRAY)) {
+      rb_raise(rb_eArgError, "last_n_tokens must be an Array");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "penalty must be a float");
+      return Qnil;
+    }
+
+    const size_t last_tokens_size = RARRAY_LEN(last_n_tokens);
+    std::vector<llama_token> last_n_tokens_data(last_tokens_size);
+    for (size_t i = 0; i < last_tokens_size; i++) {
+      last_n_tokens_data[i] = NUM2INT(rb_ary_entry(last_n_tokens, i));
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float penalty = NUM2DBL(kw_values[0]);
+
+    llama_sample_repetition_penalty(ctx_ptr->ctx, &(cnd_ptr->array), last_n_tokens_data.data(), last_tokens_size, penalty);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_frequency_and_presence_penalties(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("frequency"), rb_intern("presence") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    VALUE candidates = Qnil;
+    VALUE last_n_tokens = Qnil;
+    rb_scan_args(argc, argv, "2:", &candidates, &last_n_tokens, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "candidates must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_TYPE_P(last_n_tokens, T_ARRAY)) {
+      rb_raise(rb_eArgError, "last_n_tokens must be an Array");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "frequency must be a float");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "presence must be a float");
+      return Qnil;
+    }
+
+    const size_t last_tokens_size = RARRAY_LEN(last_n_tokens);
+    std::vector<llama_token> last_n_tokens_data(last_tokens_size);
+    for (size_t i = 0; i < last_tokens_size; i++) {
+      last_n_tokens_data[i] = NUM2INT(rb_ary_entry(last_n_tokens, i));
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+
+    const float alpha_frequency = NUM2DBL(kw_values[0]);
+    const float alpha_presence = NUM2DBL(kw_values[1]);
+
+    llama_sample_frequency_and_presence_penalties(ctx_ptr->ctx, &(cnd_ptr->array), last_n_tokens_data.data(), last_tokens_size, alpha_frequency, alpha_presence);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_softmax(VALUE self, VALUE candidates) {
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "argument must be a TokenDataArray");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+
+    llama_sample_softmax(ctx_ptr->ctx, &(cnd_ptr->array));
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_top_k(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("k"), rb_intern("min_keep") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 1, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "k must be an integer");
+      return Qnil;
+    }
+    if (kw_values[1] != Qundef && !RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "min_keep must be an integer");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const int k = NUM2DBL(kw_values[0]);
+    const size_t min_keep = kw_values[1] != Qundef ? NUM2SIZET(kw_values[1]) : 1;
+
+    llama_sample_top_k(ctx_ptr->ctx, &(cnd_ptr->array), k, min_keep);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_top_p(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("prob"), rb_intern("min_keep") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 1, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "prob must be a float");
+      return Qnil;
+    }
+    if (kw_values[1] != Qundef && !RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "min_keep must be an integer");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float prob = NUM2DBL(kw_values[0]);
+    const size_t min_keep = kw_values[1] != Qundef ? NUM2SIZET(kw_values[1]) : 1;
+
+    llama_sample_top_p(ctx_ptr->ctx, &(cnd_ptr->array), prob, min_keep);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_tail_free(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("z"), rb_intern("min_keep") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 1, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "prob must be a float");
+      return Qnil;
+    }
+    if (kw_values[1] != Qundef && !RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "min_keep must be an integer");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float z = NUM2DBL(kw_values[0]);
+    const size_t min_keep = kw_values[1] != Qundef ? NUM2SIZET(kw_values[1]) : 1;
+
+    llama_sample_tail_free(ctx_ptr->ctx, &(cnd_ptr->array), z, min_keep);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_typical(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("prob"), rb_intern("min_keep") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 1, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "prob must be a float");
+      return Qnil;
+    }
+    if (kw_values[1] != Qundef && !RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "min_keep must be an integer");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float prob = NUM2DBL(kw_values[0]);
+    const size_t min_keep = kw_values[1] != Qundef ? NUM2SIZET(kw_values[1]) : 1;
+
+    llama_sample_typical(ctx_ptr->ctx, &(cnd_ptr->array), prob, min_keep);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_temperature(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("temperature") };
+    VALUE kw_values[1] = { Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "temperature must be a float");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float temperature = NUM2DBL(kw_values[0]);
+
+    llama_sample_temperature(ctx_ptr->ctx, &(cnd_ptr->array), temperature);
+
+    return Qnil;
+  };
+
+  static VALUE _llama_context_sample_token_mirostat(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[4] = { rb_intern("tau"), rb_intern("eta"), rb_intern("m"), rb_intern("mu") };
+    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 4, 0, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "tau must be a float");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "eta must be a float");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[2])) {
+      rb_raise(rb_eArgError, "m must be an integer");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[3])) {
+      rb_raise(rb_eArgError, "mu must be a float");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float tau = NUM2DBL(kw_values[0]);
+    const float eta = NUM2DBL(kw_values[1]);
+    const int m = NUM2INT(kw_values[2]);
+    float mu = NUM2DBL(kw_values[3]);
+
+    llama_token id = llama_sample_token_mirostat(ctx_ptr->ctx, &(cnd_ptr->array), tau, eta, m, &mu);
+
+    VALUE ret = rb_ary_new2(2);
+    rb_ary_store(ret, 0, INT2NUM(id));
+    rb_ary_store(ret, 1, DBL2NUM(mu));
+    return ret;
+  };
+
+  static VALUE _llama_context_sample_token_mirostat_v2(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[3] = { rb_intern("tau"), rb_intern("eta"), rb_intern("mu") };
+    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
+    VALUE candidates = Qnil;
+    rb_scan_args(argc, argv, "1:", &candidates, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 3, 0, kw_values);
+
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "tau must be a float");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "eta must be a float");
+      return Qnil;
+    }
+    if (!RB_FLOAT_TYPE_P(kw_values[2])) {
+      rb_raise(rb_eArgError, "mu must be a float");
+      return Qnil;
+    }
+
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    const float tau = NUM2DBL(kw_values[0]);
+    const float eta = NUM2DBL(kw_values[1]);
+    float mu = NUM2DBL(kw_values[2]);
+
+    llama_token id = llama_sample_token_mirostat_v2(ctx_ptr->ctx, &(cnd_ptr->array), tau, eta, &mu);
+
+    VALUE ret = rb_ary_new2(2);
+    rb_ary_store(ret, 0, INT2NUM(id));
+    rb_ary_store(ret, 1, DBL2NUM(mu));
+    return ret;
+  };
+
+  static VALUE _llama_context_sample_token_greedy(VALUE self, VALUE candidates) {
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    llama_token id = llama_sample_token_greedy(ctx_ptr->ctx, &(cnd_ptr->array));
+    return INT2NUM(id);
+  };
+
+  static VALUE _llama_context_sample_token(VALUE self, VALUE candidates) {
+    LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
+    if (ctx_ptr->ctx == NULL) {
+      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    if (!rb_obj_is_kind_of(candidates, rb_cLLaMATokenDataArray)) {
+      rb_raise(rb_eArgError, "1st argument must be a TokenDataArray");
+      return Qnil;
+    }
+    LLaMATokenDataArrayWrapper* cnd_ptr = RbLLaMATokenDataArray::get_llama_token_data_array(candidates);
+    if (cnd_ptr->array.data == nullptr) {
+      rb_raise(rb_eRuntimeError, "TokenDataArray is empty");
+      return Qnil;
+    }
+    llama_token id = llama_sample_token(ctx_ptr->ctx, &(cnd_ptr->array));
+    return INT2NUM(id);
+  };
 };
 
 const rb_data_type_t RbLLaMAContext::llama_context_type = {
@@ -680,6 +1373,10 @@ static VALUE rb_llama_token_eos(VALUE self) {
   return INT2NUM(llama_token_eos());
 }
 
+static VALUE rb_llama_token_nl(VALUE self) {
+  return INT2NUM(llama_token_nl());
+}
+
 static VALUE rb_llama_print_system_info(VALUE self) {
   const char* result = llama_print_system_info();
   return rb_utf8_str_new_cstr(result);
@@ -695,12 +1392,16 @@ static VALUE rb_llama_mlock_supported(VALUE self) {
 
 extern "C" void Init_llama_cpp(void) {
   rb_mLLaMACpp = rb_define_module("LLaMACpp");
+
+  RbLLaMATokenData::define_class(rb_mLLaMACpp);
+  RbLLaMATokenDataArray::define_class(rb_mLLaMACpp);
   RbLLaMAContext::define_class(rb_mLLaMACpp);
   RbLLaMAContextParams::define_class(rb_mLLaMACpp);
 
   rb_define_module_function(rb_mLLaMACpp, "model_quantize", rb_llama_model_quantize, -1);
   rb_define_module_function(rb_mLLaMACpp, "token_bos", rb_llama_token_bos, 0);
   rb_define_module_function(rb_mLLaMACpp, "token_eos", rb_llama_token_eos, 0);
+  rb_define_module_function(rb_mLLaMACpp, "token_nl", rb_llama_token_nl, 0);
   rb_define_module_function(rb_mLLaMACpp, "print_system_info", rb_llama_print_system_info, 0);
   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);
@@ -711,7 +1412,6 @@ extern "C" void Init_llama_cpp(void) {
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_1", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_1));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16));
   rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_2", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_2));
-  rb_define_const(rb_mLLaMACpp, "LLAMA_FTYPE_MOSTLY_Q4_3", INT2NUM(LLAMA_FTYPE_MOSTLY_Q4_3));
   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));