llama_cpp 0.5.3 → 0.6.0

data/ext/llama_cpp/llama_cpp.cpp

@@ -1,7 +1,9 @@
 #include "llama_cpp.h"
 
 VALUE rb_mLLaMACpp;
+VALUE rb_cLLaMABatch;
 VALUE rb_cLLaMAModel;
+VALUE rb_cLLaMAModelParams;
 VALUE rb_cLLaMATimings;
 VALUE rb_cLLaMAContext;
 VALUE rb_cLLaMAContextParams;
@@ -11,6 +13,238 @@ VALUE rb_cLLaMATokenDataArray;
 VALUE rb_cLLaMAGrammarElement;
 VALUE rb_cLLaMAGrammar;
 
+class LLaMABatchWrapper {
+public:
+  llama_batch batch;
+
+  LLaMABatchWrapper() {}
+
+  ~LLaMABatchWrapper() {
+    llama_batch_free(batch);
+  }
+};
+
+class RbLLaMABatch {
+public:
+  static VALUE llama_batch_alloc(VALUE self) {
+    LLaMABatchWrapper* ptr = (LLaMABatchWrapper*)ruby_xmalloc(sizeof(LLaMABatchWrapper));
+    new (ptr) LLaMABatchWrapper();
+    return TypedData_Wrap_Struct(self, &llama_batch_type, ptr);
+  }
+
+  static void llama_batch_free(void* ptr) {
+    ((LLaMABatchWrapper*)ptr)->~LLaMABatchWrapper();
+    ruby_xfree(ptr);
+  }
+
+  static size_t llama_batch_size(const void* ptr) {
+    return sizeof(*((LLaMABatchWrapper*)ptr));
+  }
+
+  static LLaMABatchWrapper* get_llama_batch(VALUE self) {
+    LLaMABatchWrapper* ptr;
+    TypedData_Get_Struct(self, LLaMABatchWrapper, &llama_batch_type, ptr);
+    return ptr;
+  }
+
+  static void define_class(VALUE outer) {
+    rb_cLLaMABatch = rb_define_class_under(outer, "Batch", rb_cObject);
+    rb_define_alloc_func(rb_cLLaMABatch, llama_batch_alloc);
+    rb_define_method(rb_cLLaMABatch, "initialize", RUBY_METHOD_FUNC(_llama_batch_initialize), -1);
+    rb_define_method(rb_cLLaMABatch, "n_tokens=", RUBY_METHOD_FUNC(_llama_batch_set_n_tokens), 1);
+    rb_define_method(rb_cLLaMABatch, "n_tokens", RUBY_METHOD_FUNC(_llama_batch_get_n_tokens), 0);
+    rb_define_method(rb_cLLaMABatch, "all_pos_zero=", RUBY_METHOD_FUNC(_llama_batch_set_all_pos_zero), 1);
+    rb_define_method(rb_cLLaMABatch, "all_pos_zero", RUBY_METHOD_FUNC(_llama_batch_get_all_pos_zero), 0);
+    rb_define_method(rb_cLLaMABatch, "all_pos_one=", RUBY_METHOD_FUNC(_llama_batch_set_all_pos_one), 1);
+    rb_define_method(rb_cLLaMABatch, "all_pos_one", RUBY_METHOD_FUNC(_llama_batch_get_all_pos_one), 0);
+    rb_define_method(rb_cLLaMABatch, "all_seq_id=", RUBY_METHOD_FUNC(_llama_batch_set_all_seq_id), 1);
+    rb_define_method(rb_cLLaMABatch, "all_seq_id", RUBY_METHOD_FUNC(_llama_batch_get_all_seq_id), 0);
+    rb_define_method(rb_cLLaMABatch, "set_token", RUBY_METHOD_FUNC(_llama_batch_set_token), 2);
+    rb_define_method(rb_cLLaMABatch, "get_token", RUBY_METHOD_FUNC(_llama_batch_get_token), 1);
+    rb_define_method(rb_cLLaMABatch, "set_pos", RUBY_METHOD_FUNC(_llama_batch_set_pos), 2);
+    rb_define_method(rb_cLLaMABatch, "get_pos", RUBY_METHOD_FUNC(_llama_batch_get_pos), 1);
+    rb_define_method(rb_cLLaMABatch, "set_seq_id", RUBY_METHOD_FUNC(_llama_batch_set_seq_id), 2);
+    rb_define_method(rb_cLLaMABatch, "get_seq_id", RUBY_METHOD_FUNC(_llama_batch_get_seq_id), 1);
+    rb_define_method(rb_cLLaMABatch, "set_logits", RUBY_METHOD_FUNC(_llama_batch_set_logits), 2);
+    rb_define_method(rb_cLLaMABatch, "get_logits", RUBY_METHOD_FUNC(_llama_batch_get_logits), 1);
+  }
+
+private:
+  static const rb_data_type_t llama_batch_type;
+
+  static VALUE _llama_batch_initialize(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("n_tokens"), rb_intern("embd") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+    rb_scan_args(argc, argv, ":", &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "n_tokens must be an integer");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "embd must be an integer");
+      return Qnil;
+    }
+
+    const int32_t n_tokens = NUM2INT(kw_values[0]);
+    const int32_t embd = NUM2INT(kw_values[1]);
+
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    ptr->batch = llama_batch_init(n_tokens, embd);
+
+    return Qnil;
+  }
+
+  // n_tokens
+  static VALUE _llama_batch_set_n_tokens(VALUE self, VALUE n_tokens) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    ptr->batch.n_tokens = NUM2INT(n_tokens);
+    return INT2NUM(ptr->batch.n_tokens);
+  }
+
+  static VALUE _llama_batch_get_n_tokens(VALUE self) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    return INT2NUM(ptr->batch.n_tokens);
+  }
+
+  // all_pos_0
+  static VALUE _llama_batch_set_all_pos_zero(VALUE self, VALUE all_pos_0) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    ptr->batch.all_pos_0 = NUM2INT(all_pos_0);
+    return INT2NUM(ptr->batch.all_pos_0);
+  }
+
+  static VALUE _llama_batch_get_all_pos_zero(VALUE self) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    return INT2NUM(ptr->batch.all_pos_0);
+  }
+
+  // all_pos_1
+  static VALUE _llama_batch_set_all_pos_one(VALUE self, VALUE all_pos_1) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    ptr->batch.all_pos_1 = NUM2INT(all_pos_1);
+    return INT2NUM(ptr->batch.all_pos_1);
+  }
+
+  static VALUE _llama_batch_get_all_pos_one(VALUE self) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    return INT2NUM(ptr->batch.all_pos_1);
+  }
+
+  // all_seq_id
+  static VALUE _llama_batch_set_all_seq_id(VALUE self, VALUE all_seq_id) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    ptr->batch.all_seq_id = NUM2INT(all_seq_id);
+    return INT2NUM(ptr->batch.all_seq_id);
+  }
+
+  static VALUE _llama_batch_get_all_seq_id(VALUE self) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    return INT2NUM(ptr->batch.all_seq_id);
+  }
+
+  // token
+  static VALUE _llama_batch_set_token(VALUE self, VALUE idx, VALUE value) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "idx must be in [0, n_tokens)");
+      return Qnil;
+    }
+    ptr->batch.token[id] = NUM2INT(value);
+    return INT2NUM(ptr->batch.token[id]);
+  }
+
+  static VALUE _llama_batch_get_token(VALUE self, VALUE idx) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    return INT2NUM(ptr->batch.token[id]);
+  }
+
+  // pos
+  static VALUE _llama_batch_set_pos(VALUE self, VALUE idx, VALUE value) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    ptr->batch.pos[id] = NUM2INT(value);
+    return INT2NUM(ptr->batch.pos[id]);
+  }
+
+  static VALUE _llama_batch_get_pos(VALUE self, VALUE idx) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    return INT2NUM(ptr->batch.pos[id]);
+  }
+
+  // seq_id
+  static VALUE _llama_batch_set_seq_id(VALUE self, VALUE idx, VALUE value) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    ptr->batch.seq_id[id] = NUM2INT(value);
+    return INT2NUM(ptr->batch.seq_id[id]);
+  }
+
+  static VALUE _llama_batch_get_seq_id(VALUE self, VALUE idx) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    return INT2NUM(ptr->batch.seq_id[id]);
+  }
+
+  // logits
+  static VALUE _llama_batch_set_logits(VALUE self, VALUE idx, VALUE value) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    ptr->batch.logits[id] = RTEST(value) ? true : false;
+    return ptr->batch.logits[id] ? Qtrue : Qfalse;
+  }
+
+  static VALUE _llama_batch_get_logits(VALUE self, VALUE idx) {
+    LLaMABatchWrapper* ptr = get_llama_batch(self);
+    const int32_t id = NUM2INT(idx);
+    if (id < 0 || id >= ptr->batch.n_tokens) {
+      rb_raise(rb_eArgError, "id must be in [0, n_tokens)");
+      return Qnil;
+    }
+    return ptr->batch.logits[id] ? Qtrue : Qfalse;
+  }
+};
+
+const rb_data_type_t RbLLaMABatch::llama_batch_type = {
+  "RbLLaMABatch",
+  { NULL,
+    RbLLaMABatch::llama_batch_free,
+    RbLLaMABatch::llama_batch_size },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+
+};
+
 class LLaMATokenDataWrapper {
 public:
   llama_token_data data;
@@ -363,6 +597,144 @@ const rb_data_type_t RbLLaMATimings::llama_timings_type = {
   RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+class LLaMAModelParamsWrapper {
+public:
+  struct llama_model_params params;
+
+  LLaMAModelParamsWrapper() : params(llama_model_default_params()) {}
+
+  ~LLaMAModelParamsWrapper() {}
+};
+
+class RbLLaMAModelParams {
+public:
+  static VALUE llama_model_params_alloc(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = (LLaMAModelParamsWrapper*)ruby_xmalloc(sizeof(LLaMAModelParamsWrapper));
+    new (ptr) LLaMAModelParamsWrapper();
+    return TypedData_Wrap_Struct(self, &llama_model_params_type, ptr);
+  }
+
+  static void llama_model_params_free(void* ptr) {
+    ((LLaMAModelParamsWrapper*)ptr)->~LLaMAModelParamsWrapper();
+    ruby_xfree(ptr);
+  }
+
+  static size_t llama_model_params_size(const void* ptr) {
+    return sizeof(*((LLaMAModelParamsWrapper*)ptr));
+  }
+
+  static LLaMAModelParamsWrapper* get_llama_model_params(VALUE self) {
+    LLaMAModelParamsWrapper* ptr;
+    TypedData_Get_Struct(self, LLaMAModelParamsWrapper, &llama_model_params_type, ptr);
+    return ptr;
+  }
+
+  static void define_class(VALUE outer) {
+    rb_cLLaMAModelParams = rb_define_class_under(outer, "ModelParams", rb_cObject);
+    rb_define_alloc_func(rb_cLLaMAModelParams, llama_model_params_alloc);
+    rb_define_method(rb_cLLaMAModelParams, "n_gpu_layers=", RUBY_METHOD_FUNC(_llama_model_params_set_n_gpu_layers), 1);
+    rb_define_method(rb_cLLaMAModelParams, "n_gpu_layers", RUBY_METHOD_FUNC(_llama_model_params_get_n_gpu_layers), 0);
+    rb_define_method(rb_cLLaMAModelParams, "main_gpu=", RUBY_METHOD_FUNC(_llama_model_params_set_main_gpu), 1);
+    rb_define_method(rb_cLLaMAModelParams, "main_gpu", RUBY_METHOD_FUNC(_llama_model_params_get_main_gpu), 0);
+    rb_define_method(rb_cLLaMAModelParams, "tensor_split", RUBY_METHOD_FUNC(_llama_model_params_get_tensor_split), 0);
+    rb_define_method(rb_cLLaMAModelParams, "vocab_only=", RUBY_METHOD_FUNC(_llama_model_params_set_vocab_only), 1);
+    rb_define_method(rb_cLLaMAModelParams, "vocab_only", RUBY_METHOD_FUNC(_llama_model_params_get_vocab_only), 0);
+    rb_define_method(rb_cLLaMAModelParams, "use_mmap=", RUBY_METHOD_FUNC(_llama_model_params_set_use_mmap), 1);
+    rb_define_method(rb_cLLaMAModelParams, "use_mmap", RUBY_METHOD_FUNC(_llama_model_params_get_use_mmap), 0);
+    rb_define_method(rb_cLLaMAModelParams, "use_mlock=", RUBY_METHOD_FUNC(_llama_model_params_set_use_mlock), 1);
+    rb_define_method(rb_cLLaMAModelParams, "use_mlock", RUBY_METHOD_FUNC(_llama_model_params_get_use_mlock), 0);
+  }
+
+private:
+  static const rb_data_type_t llama_model_params_type;
+
+  // n_gpu_layers
+  static VALUE _llama_model_params_set_n_gpu_layers(VALUE self, VALUE n_gpu_layers) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    ptr->params.n_gpu_layers = NUM2INT(n_gpu_layers);
+    return INT2NUM(ptr->params.n_gpu_layers);
+  }
+
+  static VALUE _llama_model_params_get_n_gpu_layers(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    return INT2NUM(ptr->params.n_gpu_layers);
+  }
+
+  // main_gpu
+  static VALUE _llama_model_params_set_main_gpu(VALUE self, VALUE main_gpu) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    ptr->params.main_gpu = NUM2INT(main_gpu);
+    return INT2NUM(ptr->params.main_gpu);
+  }
+
+  static VALUE _llama_model_params_get_main_gpu(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    return INT2NUM(ptr->params.main_gpu);
+  }
+
+  // tensor_split
+  static VALUE _llama_model_params_get_tensor_split(VALUE self) {
+    if (LLAMA_MAX_DEVICES < 1) {
+      return rb_ary_new();
+    }
+    VALUE ret = rb_ary_new2(LLAMA_MAX_DEVICES);
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    if (ptr->params.tensor_split == nullptr) {
+      return rb_ary_new();
+    }
+    for (size_t i = 0; i < LLAMA_MAX_DEVICES; i++) {
+      rb_ary_store(ret, i, DBL2NUM(ptr->params.tensor_split[i]));
+    }
+    return ret;
+  }
+
+  // vocab_only
+  static VALUE _llama_model_params_set_vocab_only(VALUE self, VALUE vocab_only) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    ptr->params.vocab_only = RTEST(vocab_only) ? true : false;
+    return ptr->params.vocab_only ? Qtrue : Qfalse;
+  }
+
+  static VALUE _llama_model_params_get_vocab_only(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    return ptr->params.vocab_only ? Qtrue : Qfalse;
+  }
+
+  // use_mmap
+  static VALUE _llama_model_params_set_use_mmap(VALUE self, VALUE use_mmap) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    ptr->params.use_mmap = RTEST(use_mmap) ? true : false;
+    return ptr->params.use_mmap ? Qtrue : Qfalse;
+  }
+
+  static VALUE _llama_model_params_get_use_mmap(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    return ptr->params.use_mmap ? Qtrue : Qfalse;
+  }
+
+  // use_mlock
+  static VALUE _llama_model_params_set_use_mlock(VALUE self, VALUE use_mlock) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    ptr->params.use_mlock = RTEST(use_mlock) ? true : false;
+    return ptr->params.use_mlock ? Qtrue : Qfalse;
+  }
+
+  static VALUE _llama_model_params_get_use_mlock(VALUE self) {
+    LLaMAModelParamsWrapper* ptr = get_llama_model_params(self);
+    return ptr->params.use_mlock ? Qtrue : Qfalse;
+  }
+};
+
+const rb_data_type_t RbLLaMAModelParams::llama_model_params_type = {
+  "RbLLaMAModelParams",
+  { NULL,
+    RbLLaMAModelParams::llama_model_params_free,
+    RbLLaMAModelParams::llama_model_params_size },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+};
+
 class LLaMAContextParamsWrapper {
 public:
   struct llama_context_params params;
@@ -399,35 +771,26 @@ public:
     rb_cLLaMAContextParams = rb_define_class_under(outer, "ContextParams", rb_cObject);
     rb_define_alloc_func(rb_cLLaMAContextParams, llama_context_params_alloc);
     // rb_define_method(rb_cLLaMAContextParams, "initialize", RUBY_METHOD_FUNC(_llama_context_params_init), 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, "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, "n_threads=", RUBY_METHOD_FUNC(_llama_context_params_set_n_threads), 1);
+    rb_define_method(rb_cLLaMAContextParams, "n_threads", RUBY_METHOD_FUNC(_llama_context_params_get_n_threads), 0);
+    rb_define_method(rb_cLLaMAContextParams, "n_threads_batch=", RUBY_METHOD_FUNC(_llama_context_params_set_n_threads_batch), 1);
+    rb_define_method(rb_cLLaMAContextParams, "n_threads_batch", RUBY_METHOD_FUNC(_llama_context_params_get_n_threads_batch), 0);
     rb_define_method(rb_cLLaMAContextParams, "rope_freq_base=", RUBY_METHOD_FUNC(_llama_context_params_set_rope_freq_base), 1);
     rb_define_method(rb_cLLaMAContextParams, "rope_freq_base", RUBY_METHOD_FUNC(_llama_context_params_get_rope_freq_base), 0);
     rb_define_method(rb_cLLaMAContextParams, "rope_freq_scale=", RUBY_METHOD_FUNC(_llama_context_params_set_rope_freq_scale), 1);
     rb_define_method(rb_cLLaMAContextParams, "rope_freq_scale", RUBY_METHOD_FUNC(_llama_context_params_get_rope_freq_scale), 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, "mul_mat_q=", RUBY_METHOD_FUNC(_llama_context_params_set_mul_mat_q), 1);
     rb_define_method(rb_cLLaMAContextParams, "mul_mat_q", RUBY_METHOD_FUNC(_llama_context_params_get_mul_mat_q), 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);
     rb_define_method(rb_cLLaMAContextParams, "f16_kv", RUBY_METHOD_FUNC(_llama_context_params_get_f16_kv), 0);
     rb_define_method(rb_cLLaMAContextParams, "logits_all=", RUBY_METHOD_FUNC(_llama_context_params_set_logits_all), 1);
     rb_define_method(rb_cLLaMAContextParams, "logits_all", RUBY_METHOD_FUNC(_llama_context_params_get_logits_all), 0);
-    rb_define_method(rb_cLLaMAContextParams, "vocab_only=", RUBY_METHOD_FUNC(_llama_context_params_set_vocab_only), 1);
-    rb_define_method(rb_cLLaMAContextParams, "vocab_only", RUBY_METHOD_FUNC(_llama_context_params_get_vocab_only), 0);
-    rb_define_method(rb_cLLaMAContextParams, "use_mmap=", RUBY_METHOD_FUNC(_llama_context_params_set_use_mmap), 1);
-    rb_define_method(rb_cLLaMAContextParams, "use_mmap", RUBY_METHOD_FUNC(_llama_context_params_get_use_mmap), 0);
-    rb_define_method(rb_cLLaMAContextParams, "use_mlock=", RUBY_METHOD_FUNC(_llama_context_params_set_use_mlock), 1);
-    rb_define_method(rb_cLLaMAContextParams, "use_mlock", RUBY_METHOD_FUNC(_llama_context_params_get_use_mlock), 0);
     rb_define_method(rb_cLLaMAContextParams, "embedding=", RUBY_METHOD_FUNC(_llama_context_params_set_embedding), 1);
     rb_define_method(rb_cLLaMAContextParams, "embedding", RUBY_METHOD_FUNC(_llama_context_params_get_embedding), 0);
   }
@@ -441,6 +804,22 @@ private:
   //   return self;
   // }
 
+  // seed
+  static VALUE _llama_context_params_set_seed(VALUE self, VALUE seed) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    if (NUM2INT(seed) < 0) {
+      rb_raise(rb_eArgError, "seed must be positive");
+      return Qnil;
+    }
+    ptr->params.seed = NUM2INT(seed);
+    return INT2NUM(ptr->params.seed);
+  }
+
+  static VALUE _llama_context_params_get_seed(VALUE self) {
+    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
+    return INT2NUM(ptr->params.seed);
+  }
+
   // n_ctx
   static VALUE _llama_context_params_set_n_ctx(VALUE self, VALUE n_ctx) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
@@ -465,41 +844,28 @@ private:
     return INT2NUM(ptr->params.n_batch);
   }
 
-  // n_gpu_layers
-  static VALUE _llama_context_params_set_n_gpu_layers(VALUE self, VALUE n_gpu_layers) {
+  // n_threads
+  static VALUE _llama_context_params_set_n_threads(VALUE self, VALUE n_threads) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    ptr->params.n_gpu_layers = NUM2INT(n_gpu_layers);
-    return INT2NUM(ptr->params.n_gpu_layers);
+    ptr->params.n_threads = NUM2INT(n_threads);
+    return INT2NUM(ptr->params.n_threads);
   }
 
-  static VALUE _llama_context_params_get_n_gpu_layers(VALUE self) {
+  static VALUE _llama_context_params_get_n_threads(VALUE self) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    return INT2NUM(ptr->params.n_gpu_layers);
+    return INT2NUM(ptr->params.n_threads);
   }
 
-  // main_gpu
-  static VALUE _llama_context_params_set_main_gpu(VALUE self, VALUE main_gpu) {
+  // n_threads_batch
+  static VALUE _llama_context_params_set_n_threads_batch(VALUE self, VALUE n_threads_batch) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    ptr->params.main_gpu = NUM2INT(main_gpu);
-    return INT2NUM(ptr->params.main_gpu);
+    ptr->params.n_threads_batch = NUM2INT(n_threads_batch);
+    return INT2NUM(ptr->params.n_threads_batch);
   }
 
-  static VALUE _llama_context_params_get_main_gpu(VALUE self) {
+  static VALUE _llama_context_params_get_n_threads_batch(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;
+    return INT2NUM(ptr->params.n_threads_batch);
   }
 
   // rope_freq_base
@@ -526,18 +892,6 @@ private:
     return DBL2NUM(ptr->params.rope_freq_scale);
   }
 
-  // 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 = RTEST(low_vram) ? 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;
-  }
-
   // mul_mat_q
   static VALUE _llama_context_params_set_mul_mat_q(VALUE self, VALUE mul_mat_q) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
@@ -550,22 +904,6 @@ private:
     return ptr->params.mul_mat_q ? Qtrue : Qfalse;
   }
 
-  // seed
-  static VALUE _llama_context_params_set_seed(VALUE self, VALUE seed) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    if (NUM2INT(seed) < 0) {
-      rb_raise(rb_eArgError, "seed must be positive");
-      return Qnil;
-    }
-    ptr->params.seed = NUM2INT(seed);
-    return INT2NUM(ptr->params.seed);
-  }
-
-  static VALUE _llama_context_params_get_seed(VALUE self) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    return INT2NUM(ptr->params.seed);
-  }
-
   // f16_kv
   static VALUE _llama_context_params_set_f16_kv(VALUE self, VALUE f16_kv) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
@@ -590,42 +928,6 @@ private:
     return ptr->params.logits_all ? Qtrue : Qfalse;
   }
 
-  // vocab_only
-  static VALUE _llama_context_params_set_vocab_only(VALUE self, VALUE vocab_only) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    ptr->params.vocab_only = RTEST(vocab_only) ? true : false;
-    return ptr->params.vocab_only ? Qtrue : Qfalse;
-  }
-
-  static VALUE _llama_context_params_get_vocab_only(VALUE self) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    return ptr->params.vocab_only ? Qtrue : Qfalse;
-  }
-
-  // use_mmap
-  static VALUE _llama_context_params_set_use_mmap(VALUE self, VALUE use_mmap) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    ptr->params.use_mmap = RTEST(use_mmap) ? true : false;
-    return ptr->params.use_mmap ? Qtrue : Qfalse;
-  }
-
-  static VALUE _llama_context_params_get_use_mmap(VALUE self) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    return ptr->params.use_mmap ? Qtrue : Qfalse;
-  }
-
-  // use_mlock
-  static VALUE _llama_context_params_set_use_mlock(VALUE self, VALUE use_mlock) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    ptr->params.use_mlock = RTEST(use_mlock) ? true : false;
-    return ptr->params.use_mlock ? Qtrue : Qfalse;
-  }
-
-  static VALUE _llama_context_params_get_use_mlock(VALUE self) {
-    LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
-    return ptr->params.use_mlock ? Qtrue : Qfalse;
-  }
-
   // embedding
   static VALUE _llama_context_params_set_embedding(VALUE self, VALUE embedding) {
     LLaMAContextParamsWrapper* ptr = get_llama_context_params(self);
@@ -823,11 +1125,10 @@ public:
     rb_define_method(rb_cLLaMAModel, "load", RUBY_METHOD_FUNC(_llama_model_load), -1);
     rb_define_method(rb_cLLaMAModel, "apply_lora_from_file", RUBY_METHOD_FUNC(_llama_model_apply_lora_from_file), -1);
     rb_define_method(rb_cLLaMAModel, "n_vocab", RUBY_METHOD_FUNC(_llama_model_get_model_n_vocab), 0);
-    rb_define_method(rb_cLLaMAModel, "n_ctx", RUBY_METHOD_FUNC(_llama_model_get_model_n_ctx), 0);
     rb_define_method(rb_cLLaMAModel, "n_ctx_train", RUBY_METHOD_FUNC(_llama_model_get_model_n_ctx_train), 0);
     rb_define_method(rb_cLLaMAModel, "n_embd", RUBY_METHOD_FUNC(_llama_model_get_model_n_embd), 0);
-    rb_define_method(rb_cLLaMAModel, "token_to_piece", RUBY_METHOD_FUNC(_llama_model_token_to_piece_with_model), 1);
-    rb_define_method(rb_cLLaMAModel, "tokenize", RUBY_METHOD_FUNC(_llama_model_tokenize_with_model), -1);
+    rb_define_method(rb_cLLaMAModel, "token_to_piece", RUBY_METHOD_FUNC(_llama_model_token_to_piece), 1);
+    rb_define_method(rb_cLLaMAModel, "tokenize", RUBY_METHOD_FUNC(_llama_model_tokenize), -1);
     rb_define_method(rb_cLLaMAModel, "desc", RUBY_METHOD_FUNC(_llama_model_get_model_desc), 0);
     rb_define_method(rb_cLLaMAModel, "size", RUBY_METHOD_FUNC(_llama_model_get_model_size), 0);
     rb_define_method(rb_cLLaMAModel, "n_params", RUBY_METHOD_FUNC(_llama_model_get_model_n_params), 0);
@@ -841,30 +1142,21 @@ private:
     ID kw_table[2] = { rb_intern("model_path"), rb_intern("params") };
     VALUE kw_values[2] = { Qundef, Qundef };
     rb_scan_args(argc, argv, ":", &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 0, 2, kw_values);
-
-    if (kw_values[0] == Qundef && kw_values[1] == Qundef) {
-      rb_iv_set(self, "@params", Qnil);
-      return Qnil;
-    }
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
 
     if (!RB_TYPE_P(kw_values[0], T_STRING)) {
       rb_raise(rb_eArgError, "model_path must be a string");
       return Qnil;
     }
-    if (!rb_obj_is_kind_of(kw_values[1], rb_cLLaMAContextParams)) {
-      rb_raise(rb_eArgError, "params must be a ContextParams");
+    if (!rb_obj_is_kind_of(kw_values[1], rb_cLLaMAModelParams)) {
+      rb_raise(rb_eArgError, "params must be a ModelParams");
       return Qnil;
     }
 
     VALUE filename = kw_values[0];
-    LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(kw_values[1]);
+    LLaMAModelParamsWrapper* prms_ptr = RbLLaMAModelParams::get_llama_model_params(kw_values[1]);
     LLaMAModelWrapper* model_ptr = get_llama_model(self);
 
-    if (prms_ptr->params.seed == LLAMA_DEFAULT_SEED) {
-      prms_ptr->params.seed = time(NULL);
-    }
-
     try {
       model_ptr->model = llama_load_model_from_file(StringValueCStr(filename), prms_ptr->params);
     } catch (const std::runtime_error& e) {
@@ -912,8 +1204,8 @@ private:
       rb_raise(rb_eArgError, "model_path must be a string");
       return Qnil;
     }
-    if (!rb_obj_is_kind_of(kw_values[1], rb_cLLaMAContextParams)) {
-      rb_raise(rb_eArgError, "params must be a LLaMAContextParams");
+    if (!rb_obj_is_kind_of(kw_values[1], rb_cLLaMAModelParams)) {
+      rb_raise(rb_eArgError, "params must be a LLaMAModelParams");
       return Qnil;
     }
 
@@ -924,7 +1216,7 @@ private:
     }
 
     VALUE filename = kw_values[0];
-    LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(kw_values[1]);
+    LLaMAModelParamsWrapper* prms_ptr = RbLLaMAModelParams::get_llama_model_params(kw_values[1]);
 
     try {
       model_ptr->model = llama_load_model_from_file(StringValueCStr(filename), prms_ptr->params);
@@ -946,10 +1238,10 @@ private:
 
   static VALUE _llama_model_apply_lora_from_file(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[3] = { rb_intern("lora_path"), rb_intern("base_model_path"), rb_intern("n_threads") };
-    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
+    ID kw_table[4] = { rb_intern("lora_path"), rb_intern("base_model_path"), rb_intern("n_threads"), rb_intern("scale") };
+    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
     rb_scan_args(argc, argv, ":", &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 1, 2, kw_values);
+    rb_get_kwargs(kw_args, kw_table, 1, 3, kw_values);
 
     if (!RB_TYPE_P(kw_values[0], T_STRING)) {
       rb_raise(rb_eArgError, "lora_path must be a string");
@@ -963,13 +1255,18 @@ private:
       rb_raise(rb_eArgError, "n_threads must be an integer");
       return Qnil;
     }
+    if (kw_values[3] != Qundef && !RB_FLOAT_TYPE_P(kw_values[3])) {
+      rb_raise(rb_eArgError, "scale must be a float");
+      return Qnil;
+    }
 
     const char* lora_path = StringValueCStr(kw_values[0]);
     const char* base_model_path = kw_values[1] == Qundef ? NULL : StringValueCStr(kw_values[1]);
     const int n_threads = kw_values[2] == Qundef ? 1 : NUM2INT(kw_values[2]);
+    const float scale = kw_values[3] == Qundef ? 1.0 : NUM2DBL(kw_values[3]);
 
     LLaMAModelWrapper* ptr = get_llama_model(self);
-    if (llama_model_apply_lora_from_file(ptr->model, lora_path, base_model_path, n_threads) != 0) {
+    if (llama_model_apply_lora_from_file(ptr->model, lora_path, scale, base_model_path, n_threads) != 0) {
       rb_raise(rb_eRuntimeError, "Failed to apply LoRA");
       return Qnil;
     }
@@ -978,25 +1275,20 @@ private:
 
   static VALUE _llama_model_get_model_n_vocab(VALUE self) {
     LLaMAModelWrapper* ptr = get_llama_model(self);
-    return INT2NUM(llama_model_n_vocab(ptr->model));
-  }
-
-  static VALUE _llama_model_get_model_n_ctx(VALUE self) {
-    LLaMAModelWrapper* ptr = get_llama_model(self);
-    return INT2NUM(llama_model_n_ctx(ptr->model));
+    return INT2NUM(llama_n_vocab(ptr->model));
   }
 
   static VALUE _llama_model_get_model_n_ctx_train(VALUE self) {
     LLaMAModelWrapper* ptr = get_llama_model(self);
-    return INT2NUM(llama_model_n_ctx_train(ptr->model));
+    return INT2NUM(llama_n_ctx_train(ptr->model));
   }
 
   static VALUE _llama_model_get_model_n_embd(VALUE self) {
     LLaMAModelWrapper* ptr = get_llama_model(self);
-    return INT2NUM(llama_model_n_embd(ptr->model));
+    return INT2NUM(llama_n_embd(ptr->model));
   }
 
-  static VALUE _llama_model_token_to_piece_with_model(VALUE self, VALUE token_) {
+  static VALUE _llama_model_token_to_piece(VALUE self, VALUE token_) {
     if (!RB_INTEGER_TYPE_P(token_)) {
       rb_raise(rb_eArgError, "token must be an integer");
       return Qnil;
@@ -1004,10 +1296,10 @@ private:
     const llama_token token = NUM2INT(token_);
     LLaMAModelWrapper* ptr = get_llama_model(self);
     std::vector<char> result(8, 0);
-    const int n_tokens = llama_token_to_piece_with_model(ptr->model, token, result.data(), result.size());
+    const int n_tokens = llama_token_to_piece(ptr->model, token, result.data(), result.size());
     if (n_tokens < 0) {
       result.resize(-n_tokens);
-      const int check = llama_token_to_piece_with_model(ptr->model, token, result.data(), result.size());
+      const int check = llama_token_to_piece(ptr->model, token, result.data(), result.size());
       if (check != -n_tokens) {
         rb_raise(rb_eRuntimeError, "failed to convert");
         return Qnil;
@@ -1019,7 +1311,7 @@ private:
     return rb_str_new_cstr(ret.c_str());
   }
 
-  static VALUE _llama_model_tokenize_with_model(int argc, VALUE* argv, VALUE self) {
+  static VALUE _llama_model_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") };
     VALUE kw_values[3] = { Qundef, Qundef, Qundef };
@@ -1046,7 +1338,7 @@ private:
 
     llama_token* tokens = ALLOCA_N(llama_token, n_max_tokens);
     LLaMAModelWrapper* ptr = get_llama_model(self);
-    const int n_tokens = llama_tokenize_with_model(ptr->model, text.c_str(), text.size(), tokens, n_max_tokens, add_bos);
+    const int n_tokens = llama_tokenize(ptr->model, text.c_str(), text.size(), tokens, n_max_tokens, add_bos);
 
     if (n_tokens < 0) {
       rb_raise(rb_eRuntimeError, "failed to tokenize. The numebr of tokens (%d) is greater than n_max_tokens.", -n_tokens);
@@ -1345,11 +1637,11 @@ public:
   static void define_class(VALUE outer) {
     rb_cLLaMAContext = rb_define_class_under(outer, "Context", rb_cObject);
     rb_define_alloc_func(rb_cLLaMAContext, llama_context_alloc);
+    rb_define_attr(rb_cLLaMAContext, "model", 1, 0);
     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_embd", RUBY_METHOD_FUNC(_llama_context_eval_embd), -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, "decode", RUBY_METHOD_FUNC(_llama_context_decode), 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, "text", RUBY_METHOD_FUNC(_llama_context_text), 1);
@@ -1358,15 +1650,16 @@ public:
     rb_define_method(rb_cLLaMAContext, "token_bos", RUBY_METHOD_FUNC(_llama_context_token_bos), 0);
     rb_define_method(rb_cLLaMAContext, "token_eos", RUBY_METHOD_FUNC(_llama_context_token_eos), 0);
     rb_define_method(rb_cLLaMAContext, "token_nl", RUBY_METHOD_FUNC(_llama_context_token_nl), 0);
-    rb_define_method(rb_cLLaMAContext, "token_to_piece", RUBY_METHOD_FUNC(_llama_context_token_to_piece), 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_ctx_train", RUBY_METHOD_FUNC(_llama_context_n_ctx_train), 0);
-    rb_define_method(rb_cLLaMAContext, "n_embd", RUBY_METHOD_FUNC(_llama_context_n_embd), 0);
     rb_define_method(rb_cLLaMAContext, "timings", RUBY_METHOD_FUNC(_llama_context_get_timings), 0);
     rb_define_method(rb_cLLaMAContext, "print_timings", RUBY_METHOD_FUNC(_llama_context_print_timings), 0);
     rb_define_method(rb_cLLaMAContext, "reset_timings", RUBY_METHOD_FUNC(_llama_context_reset_timings), 0);
     rb_define_method(rb_cLLaMAContext, "kv_cache_token_count", RUBY_METHOD_FUNC(_llama_context_kv_cache_token_count), 0);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_tokens_rm", RUBY_METHOD_FUNC(_llama_context_kv_cache_tokens_rm), 2);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_seq_rm", RUBY_METHOD_FUNC(_llama_context_kv_cache_seq_rm), 3);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_seq_cp", RUBY_METHOD_FUNC(_llama_context_kv_cache_seq_cp), 4);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_seq_keep", RUBY_METHOD_FUNC(_llama_context_kv_cache_seq_keep), 1);
+    rb_define_method(rb_cLLaMAContext, "kv_cache_seq_shift", RUBY_METHOD_FUNC(_llama_context_kv_cache_seq_shift), 4);
     rb_define_method(rb_cLLaMAContext, "set_rng_seed", RUBY_METHOD_FUNC(_llama_context_set_rng_seed), 1);
     rb_define_method(rb_cLLaMAContext, "load_session_file", RUBY_METHOD_FUNC(_llama_context_load_session_file), -1);
     rb_define_method(rb_cLLaMAContext, "save_session_file", RUBY_METHOD_FUNC(_llama_context_save_session_file), -1);
@@ -1378,6 +1671,7 @@ public:
     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_temp", RUBY_METHOD_FUNC(_llama_context_sample_temp), -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);
@@ -1392,24 +1686,27 @@ private:
 
   static VALUE _llama_context_initialize(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[1] = { rb_intern("model") };
-    VALUE kw_values[1] = { Qundef };
+    ID kw_table[2] = { rb_intern("model"), rb_intern("params") };
+    VALUE kw_values[2] = { Qundef, Qundef };
     rb_scan_args(argc, argv, ":", &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
 
     VALUE model = kw_values[0];
     if (!rb_obj_is_kind_of(model, rb_cLLaMAModel)) {
       rb_raise(rb_eArgError, "model must be a Model");
       return Qnil;
     }
+    VALUE params = kw_values[1];
+    if (!rb_obj_is_kind_of(params, rb_cLLaMAContextParams)) {
+      rb_raise(rb_eArgError, "params must be a ContextParams");
+      return Qnil;
+    }
 
     LLaMAModelWrapper* model_ptr = RbLLaMAModel::get_llama_model(model);
     if (model_ptr->model == NULL) {
       rb_raise(rb_eRuntimeError, "Model is empty");
       return Qnil;
     }
-
-    VALUE params = rb_iv_get(model, "@params");
     LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(params);
     LLaMAContextWrapper* ctx_ptr = get_llama_context(self);
 
@@ -1421,6 +1718,7 @@ private:
     }
 
     rb_iv_set(self, "@model", model);
+    rb_iv_set(self, "@params", params);
     rb_iv_set(self, "@has_evaluated", Qfalse);
 
     return Qnil;
@@ -1428,8 +1726,8 @@ private:
 
   static VALUE _llama_context_eval(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[4] = { rb_intern("tokens"), rb_intern("n_past"), rb_intern("n_tokens"), rb_intern("n_threads") };
-    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
+    ID kw_table[3] = { rb_intern("tokens"), rb_intern("n_past"), rb_intern("n_tokens") };
+    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
     rb_scan_args(argc, argv, ":", &kw_args);
     rb_get_kwargs(kw_args, kw_table, 2, 2, kw_values);
 
@@ -1445,10 +1743,6 @@ private:
       rb_raise(rb_eArgError, "n_tokens 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");
-      return Qnil;
-    }
 
     const size_t tokens_len = RARRAY_LEN(kw_values[0]);
     std::vector<llama_token> embd(tokens_len);
@@ -1463,14 +1757,13 @@ private:
 
     const int n_tokens = kw_values[2] == Qundef ? (int)tokens_len : NUM2INT(kw_values[2]);
     const int n_past = NUM2INT(kw_values[1]);
-    const int n_threads = kw_values[3] == Qundef ? 1 : NUM2INT(kw_values[3]);
 
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
       rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
       return Qnil;
     }
-    if (llama_eval(ptr->ctx, embd.data(), n_tokens, n_past, n_threads) != 0) {
+    if (llama_eval(ptr->ctx, embd.data(), n_tokens, n_past) != 0) {
       rb_raise(rb_eRuntimeError, "Failed to evaluate");
       return Qnil;
     }
@@ -1483,8 +1776,8 @@ private:
 
   static VALUE _llama_context_eval_embd(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[4] = { rb_intern("embd"), rb_intern("n_past"), rb_intern("n_tokens"), rb_intern("n_threads") };
-    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
+    ID kw_table[3] = { rb_intern("embd"), rb_intern("n_past"), rb_intern("n_tokens") };
+    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
     rb_scan_args(argc, argv, ":", &kw_args);
     rb_get_kwargs(kw_args, kw_table, 2, 2, kw_values);
 
@@ -1500,10 +1793,6 @@ private:
       rb_raise(rb_eArgError, "n_tokens 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");
-      return Qnil;
-    }
 
     const size_t tokens_len = RARRAY_LEN(kw_values[0]);
     std::vector<float> embd(tokens_len);
@@ -1518,14 +1807,13 @@ private:
 
     const int n_tokens = kw_values[2] == Qundef ? (int)tokens_len : NUM2INT(kw_values[2]);
     const int n_past = NUM2INT(kw_values[1]);
-    const int n_threads = kw_values[3] == Qundef ? 1 : NUM2INT(kw_values[3]);
 
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
       rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
       return Qnil;
     }
-    if (llama_eval_embd(ptr->ctx, embd.data(), n_tokens, n_past, n_threads) != 0) {
+    if (llama_eval_embd(ptr->ctx, embd.data(), n_tokens, n_past) != 0) {
       rb_raise(rb_eRuntimeError, "Failed to evaluate");
       return Qnil;
     }
@@ -1536,91 +1824,22 @@ private:
     return Qnil;
   }
 
-  static VALUE _llama_context_eval_export(VALUE self, VALUE fname_) {
+  static VALUE _llama_context_decode(VALUE self, VALUE batch) {
     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");
+    if (!rb_obj_is_kind_of(batch, rb_cLLaMABatch)) {
+      rb_raise(rb_eArgError, "batch must be a Batch");
       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") };
-    VALUE kw_values[3] = { Qundef, Qundef, Qundef };
-    rb_scan_args(argc, argv, ":", &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 1, 2, kw_values);
-
-    if (!RB_TYPE_P(kw_values[0], T_STRING)) {
-      rb_raise(rb_eArgError, "text must be a String");
+    LLaMABatchWrapper* batch_ptr = RbLLaMABatch::get_llama_batch(batch);
+    if (llama_decode(ptr->ctx, batch_ptr->batch) < 0) {
+      rb_raise(rb_eRuntimeError, "Failed to decode");
       return Qnil;
     }
-    if (kw_values[1] != Qundef && !RB_INTEGER_TYPE_P(kw_values[1])) {
-      rb_raise(rb_eArgError, "n_max_tokens must be an integer");
-      return Qnil;
-    }
-    if (kw_values[2] != Qundef && (kw_values[2] != Qtrue && kw_values[2] != Qfalse)) {
-      rb_raise(rb_eArgError, "add_bos must be a boolean");
-      return Qnil;
-    }
-
-    VALUE text_ = kw_values[0];
-    std::string text = StringValueCStr(text_);
-    const bool add_bos = kw_values[2] == Qtrue ? true : false;
-    const int n_max_tokens = kw_values[1] != Qundef ? NUM2INT(kw_values[1]) : text.size() + (add_bos ? 1 : 0);
-
-    std::vector<llama_token> tokens(n_max_tokens);
-    LLaMAContextWrapper* ptr = get_llama_context(self);
-    if (ptr->ctx == NULL) {
-      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
-      return Qnil;
-    }
-    const int n = llama_tokenize(ptr->ctx, text.c_str(), text.size(), tokens.data(), n_max_tokens, add_bos);
-    if (n < 0) {
-      rb_raise(rb_eRuntimeError, "Failed to tokenize");
-      return Qnil;
-    }
-
-    VALUE output = rb_ary_new();
-    for (int i = 0; i < n; i++) {
-      rb_ary_push(output, INT2NUM(tokens[i]));
-    }
-
-    RB_GC_GUARD(text_);
-    return output;
-  }
-
-  static VALUE _llama_context_token_to_piece(VALUE self, VALUE token_) {
-    LLaMAContextWrapper* ptr = get_llama_context(self);
-    if (ptr->ctx == NULL) {
-      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
-      return Qnil;
-    }
-    const llama_token token = NUM2INT(token_);
-    std::vector<char> result(8, 0);
-    const int n_tokens = llama_token_to_piece(ptr->ctx, token, result.data(), result.size());
-    if (n_tokens < 0) {
-      result.resize(-n_tokens);
-      const int check = llama_token_to_piece(ptr->ctx, token, result.data(), result.size());
-      if (check != -n_tokens) {
-        rb_raise(rb_eRuntimeError, "failed to convert");
-        return Qnil;
-      }
-    } else {
-      result.resize(n_tokens);
-    }
-    std::string ret(result.data(), result.size());
-    return rb_str_new_cstr(ret.c_str());
+    return Qnil;
   }
 
   static VALUE _llama_context_logits(VALUE self) {
@@ -1635,10 +1854,11 @@ private:
     }
 
     VALUE model = rb_iv_get(self, "@model");
-    VALUE params = rb_iv_get(model, "@params");
+    LLaMAModelWrapper* model_ptr = RbLLaMAModel::get_llama_model(model);
+    VALUE params = rb_iv_get(self, "@params");
     LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(params);
     const int n_tokens = prms_ptr->params.logits_all ? NUM2INT(rb_iv_get(self, "@n_tokens")) : 1;
-    const int n_vocab = llama_n_vocab(ptr->ctx);
+    const int n_vocab = llama_n_vocab(model_ptr->model);
     const float* logits = llama_get_logits(ptr->ctx);
     VALUE output = rb_ary_new();
     for (int i = 0; i < n_tokens * n_vocab; i++) {
@@ -1655,7 +1875,8 @@ private:
       return Qnil;
     }
     VALUE model = rb_iv_get(self, "@model");
-    VALUE params = rb_iv_get(model, "@params");
+    LLaMAModelWrapper* model_ptr = RbLLaMAModel::get_llama_model(model);
+    VALUE params = rb_iv_get(self, "@params");
     LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(params);
     if (!prms_ptr->params.embedding) {
       rb_raise(rb_eRuntimeError, "embedding parameter is false");
@@ -1666,7 +1887,7 @@ private:
       return Qnil;
     }
 
-    const int n_embd = llama_n_embd(ptr->ctx);
+    const int n_embd = llama_n_embd(model_ptr->model);
     const float* embd = llama_get_embeddings(ptr->ctx);
     VALUE output = rb_ary_new();
     for (int i = 0; i < n_embd; i++) {
@@ -1736,81 +1957,104 @@ private:
     return INT2NUM(llama_token_nl(ptr->ctx));
   }
 
-  static VALUE _llama_context_n_vocab(VALUE self) {
+  static VALUE _llama_context_n_ctx(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_n_vocab(ptr->ctx));
+    return INT2NUM(llama_n_ctx(ptr->ctx));
   }
 
-  static VALUE _llama_context_n_ctx(VALUE self) {
+  static VALUE _llama_context_get_timings(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_n_ctx(ptr->ctx));
+    VALUE tm_obj = rb_funcall(rb_cLLaMATimings, rb_intern("new"), 0);
+    LLaMATimingsWrapper* tm_ptr = RbLLaMATimings::get_llama_timings(tm_obj);
+    tm_ptr->timings = llama_get_timings(ptr->ctx);
+    return tm_obj;
   }
 
-  static VALUE _llama_context_n_ctx_train(VALUE self) {
+  static VALUE _llama_context_print_timings(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_n_ctx_train(ptr->ctx));
+    llama_print_timings(ptr->ctx);
+    return Qnil;
   }
 
-  static VALUE _llama_context_n_embd(VALUE self) {
+  static VALUE _llama_context_reset_timings(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_n_embd(ptr->ctx));
+    llama_reset_timings(ptr->ctx);
+    return Qnil;
   }
 
-  static VALUE _llama_context_get_timings(VALUE self) {
+  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;
     }
-    VALUE tm_obj = rb_funcall(rb_cLLaMATimings, rb_intern("new"), 0);
-    LLaMATimingsWrapper* tm_ptr = RbLLaMATimings::get_llama_timings(tm_obj);
-    tm_ptr->timings = llama_get_timings(ptr->ctx);
-    return tm_obj;
+    return INT2NUM(llama_get_kv_cache_token_count(ptr->ctx));
   }
 
-  static VALUE _llama_context_print_timings(VALUE self) {
+  static VALUE _llama_context_kv_cache_tokens_rm(VALUE self, VALUE c0, VALUE c1) {
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
       rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
       return Qnil;
     }
-    llama_print_timings(ptr->ctx);
+    llama_kv_cache_tokens_rm(ptr->ctx, NUM2INT(c0), NUM2INT(c1));
     return Qnil;
   }
 
-  static VALUE _llama_context_reset_timings(VALUE self) {
+  static VALUE _llama_context_kv_cache_seq_rm(VALUE self, VALUE seq_id, VALUE p0, VALUE p1) {
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
       rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
       return Qnil;
     }
-    llama_reset_timings(ptr->ctx);
+    llama_kv_cache_seq_rm(ptr->ctx, NUM2INT(seq_id), NUM2INT(p0), NUM2INT(p1));
     return Qnil;
   }
 
-  static VALUE _llama_context_kv_cache_token_count(VALUE self) {
+  static VALUE _llama_context_kv_cache_seq_cp(VALUE self, VALUE seq_id_src, VALUE seq_id_dst, VALUE p0, VALUE p1) {
     LLaMAContextWrapper* ptr = get_llama_context(self);
     if (ptr->ctx == NULL) {
-      rb_raise(rb_eRuntimeError, "LLaMA context is not initialized");
+      rb_raise(rb_eArgError, "LLaMA context is not initialized");
       return Qnil;
     }
-    return INT2NUM(llama_get_kv_cache_token_count(ptr->ctx));
+    llama_kv_cache_seq_cp(ptr->ctx, NUM2INT(seq_id_src), NUM2INT(seq_id_dst), NUM2INT(p0), NUM2INT(p1));
+    return Qnil;
+  }
+
+  static VALUE _llama_context_kv_cache_seq_keep(VALUE self, VALUE seq_id) {
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eArgError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    llama_kv_cache_seq_keep(ptr->ctx, NUM2INT(seq_id));
+    return Qnil;
+  }
+
+  static VALUE _llama_context_kv_cache_seq_shift(VALUE self, VALUE seq_id, VALUE p0, VALUE p1, VALUE delta) {
+    LLaMAContextWrapper* ptr = get_llama_context(self);
+    if (ptr->ctx == NULL) {
+      rb_raise(rb_eArgError, "LLaMA context is not initialized");
+      return Qnil;
+    }
+    llama_kv_cache_seq_shift(ptr->ctx, NUM2INT(seq_id), NUM2INT(p0), NUM2INT(p1), NUM2INT(delta));
+    return Qnil;
   }
 
   static VALUE _llama_context_set_rng_seed(VALUE self, VALUE seed_) {
@@ -1851,7 +2095,7 @@ private:
     }
 
     VALUE model = rb_iv_get(self, "@model");
-    VALUE params = rb_iv_get(model, "@params");
+    VALUE params = rb_iv_get(self, "@params");
     LLaMAContextParamsWrapper* prms_ptr = RbLLaMAContextParams::get_llama_context_params(params);
     const int n_ctx = prms_ptr->params.n_ctx;
 
@@ -2235,6 +2479,40 @@ private:
     return Qnil;
   }
 
+  static VALUE _llama_context_sample_temp(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("temp") };
+    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, "temp 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 temp = NUM2DBL(kw_values[0]);
+
+    llama_sample_temp(ctx_ptr->ctx, &(cnd_ptr->array), temp);
+
+    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") };
@@ -2560,6 +2838,7 @@ extern "C" void Init_llama_cpp(void) {
   RbLLaMATokenData::define_class(rb_mLLaMACpp);
   RbLLaMATokenDataArray::define_class(rb_mLLaMACpp);
   RbLLaMAModel::define_class(rb_mLLaMACpp);
+  RbLLaMAModelParams::define_class(rb_mLLaMACpp);
   RbLLaMATimings::define_class(rb_mLLaMACpp);
   RbLLaMAContext::define_class(rb_mLLaMACpp);
   RbLLaMAContextParams::define_class(rb_mLLaMACpp);
@@ -2578,10 +2857,6 @@ extern "C" void Init_llama_cpp(void) {
 
   rb_define_const(rb_mLLaMACpp, "LLAMA_MAX_DEVICES", INT2NUM(LLAMA_MAX_DEVICES));
 
-  rb_define_const(rb_mLLaMACpp, "LLAMA_LOG_LEVEL_ERROR", INT2NUM(LLAMA_LOG_LEVEL_ERROR));
-  rb_define_const(rb_mLLaMACpp, "LLAMA_LOG_LEVEL_WARN", INT2NUM(LLAMA_LOG_LEVEL_WARN));
-  rb_define_const(rb_mLLaMACpp, "LLAMA_LOG_LEVEL_INFO", INT2NUM(LLAMA_LOG_LEVEL_INFO));
-
   rb_define_const(rb_mLLaMACpp, "LLAMA_VOCAB_TYPE_SPM", INT2NUM(LLAMA_VOCAB_TYPE_SPM));
   rb_define_const(rb_mLLaMACpp, "LLAMA_VOCAB_TYPE_BPE", INT2NUM(LLAMA_VOCAB_TYPE_BPE));