whispercpp 1.3.5 → 1.3.6

data/ext/ruby_whisper_context.c

@@ -1,5 +1,3 @@
-#include <ruby.h>
-#include <ruby/memory_view.h>
 #include "ruby_whisper.h"
 
 extern ID id_to_s;
@@ -20,6 +18,7 @@ extern VALUE eError;
 extern VALUE cModel;
 
 extern const rb_data_type_t ruby_whisper_params_type;
+extern const rb_data_type_t ruby_whisper_context_params_type;
 extern VALUE ruby_whisper_transcribe(int argc, VALUE *argv, VALUE self);
 extern VALUE rb_whisper_model_s_new(VALUE context);
 extern VALUE rb_whisper_segment_s_new(VALUE context, int index);
@@ -27,6 +26,27 @@ extern void prepare_transcription(ruby_whisper_params *rwp, VALUE *context);
 
 ID transcribe_option_names[1];
 
+typedef struct fill_samples_args {
+  float *dest;
+  VALUE *src;
+  int n_samples;
+} fill_samples_args;
+
+typedef struct full_args {
+  VALUE *context;
+  VALUE *params;
+  float *samples;
+  int n_samples;
+} full_args;
+
+typedef struct full_parallel_args {
+  VALUE *context;
+  VALUE *params;
+  float *samples;
+  int n_samples;
+  int n_processors;
+} full_parallel_args;
+
 static void
 ruby_whisper_free(ruby_whisper *rw)
 {
@@ -124,16 +144,25 @@ ruby_whisper_initialize(int argc, VALUE *argv, VALUE self)
 {
   ruby_whisper *rw;
   VALUE whisper_model_file_path;
+  VALUE context_params;
+  struct whisper_context_params params;
 
   // TODO: we can support init from buffer here too maybe another ruby object to expose
-  rb_scan_args(argc, argv, "01", &whisper_model_file_path);
+  rb_scan_args(argc, argv, "11", &whisper_model_file_path, &context_params);
   TypedData_Get_Struct(self, ruby_whisper, &ruby_whisper_type, rw);
 
   whisper_model_file_path = ruby_whisper_normalize_model_path(whisper_model_file_path);
   if (!rb_respond_to(whisper_model_file_path, id_to_s)) {
     rb_raise(rb_eRuntimeError, "Expected file path to model to initialize Whisper::Context");
   }
-  rw->context = whisper_init_from_file_with_params(StringValueCStr(whisper_model_file_path), whisper_context_default_params());
+  if (NIL_P(context_params)) {
+    params = whisper_context_default_params();
+  } else {
+    ruby_whisper_context_params *rwcp;
+    GetContextParams(context_params, rwcp);
+    params = rwcp->params;
+  }
+  rw->context = whisper_init_from_file_with_params(StringValueCStr(whisper_model_file_path), params);
   if (rw->context == NULL) {
     rb_raise(rb_eRuntimeError, "error: failed to initialize whisper context");
   }
@@ -272,82 +301,175 @@ VALUE ruby_whisper_model_type(VALUE self)
   return rb_str_new2(whisper_model_type_readable(rw->context));
 }
 
-/*
- * Run the entire model: PCM -> log mel spectrogram -> encoder -> decoder -> text
- * Not thread safe for same context
- * Uses the specified decoding strategy to obtain the text.
- *
- * call-seq:
- *   full(params, samples, n_samples) -> nil
- *   full(params, samples) -> nil
- *
- * The second argument +samples+ must be an array of samples, respond to :length, or be a MemoryView of an array of float. It must be 32 bit float PCM audio data.
- */
-VALUE ruby_whisper_full(int argc, VALUE *argv, VALUE self)
+static bool
+check_memory_view(rb_memory_view_t *memview)
 {
-  if (argc < 2 || argc > 3) {
-    rb_raise(rb_eArgError, "wrong number of arguments (given %d, expected 2..3)", argc);
+  if (memview->format != NULL && strcmp(memview->format, "f") != 0) {
+    rb_warn("currently only format \"f\" is supported for MemoryView, but given: %s", memview->format);
+    return false;
+  }
+  if (memview->format != NULL && memview->ndim != 1) {
+    rb_warn("currently only 1 dimensional MemoryView is supported, but given: %zd", memview->ndim);
+    return false;
   }
 
-  ruby_whisper *rw;
-  ruby_whisper_params *rwp;
-  GetContext(self, rw);
-  VALUE params = argv[0];
-  TypedData_Get_Struct(params, ruby_whisper_params, &ruby_whisper_params_type, rwp);
-  VALUE samples = argv[1];
-  int n_samples;
-  rb_memory_view_t view;
-  const bool memory_view_available_p = rb_memory_view_available_p(samples);
-  if (argc == 3) {
-    n_samples = NUM2INT(argv[2]);
-    if (TYPE(samples) == T_ARRAY) {
-      if (RARRAY_LEN(samples) < n_samples) {
-        rb_raise(rb_eArgError, "samples length %ld is less than n_samples %d", RARRAY_LEN(samples), n_samples);
+  return true;
+}
+
+static VALUE
+fill_samples(VALUE rb_args)
+{
+  fill_samples_args *args = (fill_samples_args *)rb_args;
+
+  if (RB_TYPE_P(*args->src, T_ARRAY)) {
+    for (int i = 0; i < args->n_samples; i++) {
+      args->dest[i] = RFLOAT_VALUE(rb_ary_entry(*args->src, i));
+    }
+  } else {
+    // TODO: use rb_block_call
+    VALUE iter = rb_funcall(*args->src, id_to_enum, 1, rb_str_new2("each"));
+    for (int i = 0; i < args->n_samples; i++) {
+      // TODO: check if iter is exhausted and raise ArgumentError appropriately
+      VALUE sample = rb_funcall(iter, id_next, 0);
+      args->dest[i] = RFLOAT_VALUE(sample);
+    }
+  }
+
+  return Qnil;
+}
+
+struct parsed_samples_t
+parse_samples(VALUE *samples, VALUE *n_samples)
+{
+  bool memview_available = rb_memory_view_available_p(*samples);
+  struct parsed_samples_t parsed = {0};
+  parsed.memview_exported = false;
+  const bool is_array = RB_TYPE_P(*samples, T_ARRAY);
+
+  if (!NIL_P(*n_samples)) {
+    parsed.n_samples = NUM2INT(*n_samples);
+    if (is_array) {
+      if (RARRAY_LEN(*samples) < parsed.n_samples) {
+        rb_raise(rb_eArgError, "samples length %ld is less than n_samples %d", RARRAY_LEN(*samples), parsed.n_samples);
       }
     }
     // Should check when samples.respond_to?(:length)?
   } else {
-    if (TYPE(samples) == T_ARRAY) {
-      if (RARRAY_LEN(samples) > INT_MAX) {
+    if (is_array) {
+      if (RARRAY_LEN(*samples) > INT_MAX) {
         rb_raise(rb_eArgError, "samples are too long");
       }
-      n_samples = (int)RARRAY_LEN(samples);
-    } else if (memory_view_available_p) {
-      if (!rb_memory_view_get(samples, &view, RUBY_MEMORY_VIEW_SIMPLE)) {
-        view.obj = Qnil;
-        rb_raise(rb_eArgError, "unable to get a memory view");
+      parsed.n_samples = (int)RARRAY_LEN(*samples);
+    } else if (memview_available) {
+      bool memview_got = rb_memory_view_get(*samples, &parsed.memview, RUBY_MEMORY_VIEW_SIMPLE);
+      if (memview_got) {
+        parsed.memview_exported = check_memory_view(&parsed.memview);
+        if (!parsed.memview_exported) {
+          rb_memory_view_release(&parsed.memview);
+          parsed.memview = (rb_memory_view_t){0};
+        }
       }
-      ssize_t n_samples_size = view.byte_size / view.item_size;
-      if (n_samples_size > INT_MAX) {
-        rb_raise(rb_eArgError, "samples are too long");
+      if (parsed.memview_exported) {
+        ssize_t n_samples_size = parsed.memview.byte_size / parsed.memview.item_size;
+        if (n_samples_size > INT_MAX) {
+          rb_memory_view_release(&parsed.memview);
+          rb_raise(rb_eArgError, "samples are too long: %zd", n_samples_size);
+        }
+        parsed.n_samples = (int)n_samples_size;
+      } else {
+        rb_warn("unable to get a memory view. falls back to Ruby object");
+        if (rb_respond_to(*samples, id_length)) {
+          parsed.n_samples = NUM2INT(rb_funcall(*samples, id_length, 0));
+        } else {
+          rb_raise(rb_eArgError, "samples must respond to :length");
+        }
       }
-      n_samples = (int)n_samples_size;
-    } else if (rb_respond_to(samples, id_length)) {
-      n_samples = NUM2INT(rb_funcall(samples, id_length, 0));
+    } else if (rb_respond_to(*samples, id_length)) {
+      parsed.n_samples = NUM2INT(rb_funcall(*samples, id_length, 0));
     } else {
-      rb_raise(rb_eArgError, "samples must respond to :length or be a MemoryView of an array of flaot when n_samples is not given");
+      rb_raise(rb_eArgError, "samples must respond to :length or be a MemoryView of an array of float when n_samples is not given");
     }
   }
-  float * c_samples = (float *)malloc(n_samples * sizeof(float));
-  if (memory_view_available_p)  {
-    c_samples = (float *)view.data;
+
+  if (parsed.memview_exported)  {
+    parsed.samples = (float *)parsed.memview.data;
   } else {
-    if (TYPE(samples) == T_ARRAY) {
-      for (int i = 0; i < n_samples; i++) {
-        c_samples[i] = RFLOAT_VALUE(rb_ary_entry(samples, i));
-      }
-    } else {
-      // TODO: use rb_block_call
-      VALUE iter = rb_funcall(samples, id_to_enum, 1, rb_str_new2("each"));
-      for (int i = 0; i < n_samples; i++) {
-        // TODO: check if iter is exhausted and raise ArgumentError appropriately
-        VALUE sample = rb_funcall(iter, id_next, 0);
-        c_samples[i] = RFLOAT_VALUE(sample);
-      }
+    parsed.samples = ALLOC_N(float, parsed.n_samples);
+    fill_samples_args args = {
+      parsed.samples,
+      samples,
+      parsed.n_samples,
+    };
+    int state;
+    rb_protect(fill_samples, (VALUE)&args, &state);
+    if (state) {
+      xfree(parsed.samples);
+      rb_jump_tag(state);
     }
   }
-  prepare_transcription(rwp, &self);
-  const int result = whisper_full(rw->context, rwp->params, c_samples, n_samples);
+
+  return parsed;
+}
+
+VALUE
+release_samples(VALUE rb_parsed_args)
+{
+  parsed_samples_t *parsed_args = (parsed_samples_t *)rb_parsed_args;
+
+  if (parsed_args->memview_exported) {
+    rb_memory_view_release(&parsed_args->memview);
+  } else {
+    xfree(parsed_args->samples);
+  }
+  *parsed_args = (parsed_samples_t){0};
+
+  return Qnil;
+}
+
+static VALUE
+full_body(VALUE rb_args)
+{
+  full_args *args = (full_args *)rb_args;
+
+  ruby_whisper *rw;
+  ruby_whisper_params *rwp;
+  GetContext(*args->context, rw);
+  TypedData_Get_Struct(*args->params, ruby_whisper_params, &ruby_whisper_params_type, rwp);
+
+  prepare_transcription(rwp, args->context);
+  int result = whisper_full(rw->context, rwp->params, args->samples, args->n_samples);
+
+  return INT2NUM(result);
+}
+
+/*
+ * Run the entire model: PCM -> log mel spectrogram -> encoder -> decoder -> text
+ * Not thread safe for same context
+ * Uses the specified decoding strategy to obtain the text.
+ *
+ * call-seq:
+ *   full(params, samples, n_samples) -> nil
+ *   full(params, samples) -> nil
+ *
+ * The second argument +samples+ must be an array of samples, respond to :length, or be a MemoryView of an array of float. It must be 32 bit float PCM audio data.
+ */
+VALUE ruby_whisper_full(int argc, VALUE *argv, VALUE self)
+{
+  if (argc < 2 || argc > 3) {
+    rb_raise(rb_eArgError, "wrong number of arguments (given %d, expected 2..3)", argc);
+  }
+
+  VALUE n_samples = argc == 2 ? Qnil : argv[2];
+
+  struct parsed_samples_t parsed = parse_samples(&argv[1], &n_samples);
+  full_args args = {
+    &self,
+    &argv[0],
+    parsed.samples,
+    parsed.n_samples,
+  };
+  VALUE rb_result = rb_ensure(full_body, (VALUE)&args, release_samples, (VALUE)&parsed);
+  const int result = NUM2INT(rb_result);
   if (0 == result) {
     return self;
   } else {
@@ -355,6 +477,22 @@ VALUE ruby_whisper_full(int argc, VALUE *argv, VALUE self)
   }
 }
 
+static VALUE
+full_parallel_body(VALUE rb_args)
+{
+  full_parallel_args *args = (full_parallel_args *)rb_args;
+
+  ruby_whisper *rw;
+  ruby_whisper_params *rwp;
+  GetContext(*args->context, rw);
+  TypedData_Get_Struct(*args->params, ruby_whisper_params, &ruby_whisper_params_type, rwp);
+
+  prepare_transcription(rwp, args->context);
+  int result = whisper_full_parallel(rw->context, rwp->params, args->samples, args->n_samples, args->n_processors);
+
+  return INT2NUM(result);
+}
+
 /*
  * Split the input audio in chunks and process each chunk separately using whisper_full_with_state()
  * Result is stored in the default state of the context
@@ -372,19 +510,11 @@ static VALUE
 ruby_whisper_full_parallel(int argc, VALUE *argv,VALUE self)
 {
   if (argc < 2 || argc > 4) {
-    rb_raise(rb_eArgError, "wrong number of arguments (given %d, expected 2..3)", argc);
+    rb_raise(rb_eArgError, "wrong number of arguments (given %d, expected 2..4)", argc);
   }
 
-  ruby_whisper *rw;
-  ruby_whisper_params *rwp;
-  GetContext(self, rw);
-  VALUE params = argv[0];
-  TypedData_Get_Struct(params, ruby_whisper_params, &ruby_whisper_params_type, rwp);
-  VALUE samples = argv[1];
-  int n_samples;
+  VALUE n_samples = argc == 2 ? Qnil : argv[2];
   int n_processors;
-  rb_memory_view_t view;
-  const bool memory_view_available_p = rb_memory_view_available_p(samples);
   switch (argc) {
   case 2:
     n_processors = 1;
@@ -396,56 +526,16 @@ ruby_whisper_full_parallel(int argc, VALUE *argv,VALUE self)
     n_processors = NUM2INT(argv[3]);
     break;
   }
-  if (argc >= 3 && !NIL_P(argv[2])) {
-    n_samples = NUM2INT(argv[2]);
-    if (TYPE(samples) == T_ARRAY) {
-      if (RARRAY_LEN(samples) < n_samples) {
-        rb_raise(rb_eArgError, "samples length %ld is less than n_samples %d", RARRAY_LEN(samples), n_samples);
-      }
-    }
-    // Should check when samples.respond_to?(:length)?
-  } else if (memory_view_available_p) {
-    if (!rb_memory_view_get(samples, &view, RUBY_MEMORY_VIEW_SIMPLE)) {
-      view.obj = Qnil;
-      rb_raise(rb_eArgError, "unable to get a memory view");
-    }
-    ssize_t n_samples_size = view.byte_size / view.item_size;
-    if (n_samples_size > INT_MAX) {
-      rb_raise(rb_eArgError, "samples are too long");
-    }
-    n_samples = (int)n_samples_size;
-  } else {
-    if (TYPE(samples) == T_ARRAY) {
-      if (RARRAY_LEN(samples) > INT_MAX) {
-        rb_raise(rb_eArgError, "samples are too long");
-      }
-      n_samples = (int)RARRAY_LEN(samples);
-    } else if (rb_respond_to(samples, id_length)) {
-      n_samples = NUM2INT(rb_funcall(samples, id_length, 0));
-    } else {
-      rb_raise(rb_eArgError, "samples must respond to :length or be a MemoryView of an array of flaot when n_samples is not given");
-    }
-  }
-  float * c_samples = (float *)malloc(n_samples * sizeof(float));
-  if (memory_view_available_p) {
-    c_samples = (float *)view.data;
-  } else {
-    if (TYPE(samples) == T_ARRAY) {
-      for (int i = 0; i < n_samples; i++) {
-        c_samples[i] = RFLOAT_VALUE(rb_ary_entry(samples, i));
-      }
-    } else {
-      // FIXME: use rb_block_call
-      VALUE iter = rb_funcall(samples, id_to_enum, 1, rb_str_new2("each"));
-      for (int i = 0; i < n_samples; i++) {
-        // TODO: check if iter is exhausted and raise ArgumentError
-        VALUE sample = rb_funcall(iter, id_next, 0);
-        c_samples[i] = RFLOAT_VALUE(sample);
-      }
-    }
-  }
-  prepare_transcription(rwp, &self);
-  const int result = whisper_full_parallel(rw->context, rwp->params, c_samples, n_samples, n_processors);
+  struct parsed_samples_t parsed = parse_samples(&argv[1], &n_samples);
+  const full_parallel_args args = {
+    &self,
+    &argv[0],
+    parsed.samples,
+    parsed.n_samples,
+    n_processors,
+  };
+  const VALUE rb_result = rb_ensure(full_parallel_body, (VALUE)&args, release_samples, (VALUE)&parsed);
+  const int result = NUM2INT(rb_result);
   if (0 == result) {
     return self;
   } else {
@@ -631,7 +721,7 @@ ruby_whisper_get_model(VALUE self)
   return rb_whisper_model_s_new(self);
 }
 
-void
+VALUE
 init_ruby_whisper_context(VALUE *mWhisper)
 {
   cContext = rb_define_class_under(*mWhisper, "Context", rb_cObject);
@@ -669,4 +759,6 @@ init_ruby_whisper_context(VALUE *mWhisper)
   rb_define_method(cContext, "each_segment", ruby_whisper_each_segment, 0);
 
   rb_define_method(cContext, "model", ruby_whisper_get_model, 0);
+
+  return cContext;
 }

data/ext/ruby_whisper_context_params.c

@@ -0,0 +1,163 @@
+#include "ruby_whisper.h"
+
+#define NUM_PARAMS 6
+
+#define DEF_BOOLEAN_ATTR_METHOD(name) \
+static VALUE \
+ruby_whisper_context_params_get_ ## name(VALUE self) { \
+  ruby_whisper_context_params *rwcp; \
+  GetContextParams(self, rwcp); \
+  return rwcp->params.name ? Qtrue : Qfalse; \
+} \
+static VALUE \
+ruby_whisper_context_params_set_ ## name(VALUE self, VALUE value) { \
+  ruby_whisper_context_params *rwcp; \
+  GetContextParams(self, rwcp); \
+  rwcp->params.name = RTEST(value); \
+  return value; \
+}
+
+#define DEF_INT_ATTR_METHOD(name) \
+static VALUE \
+ruby_whisper_context_params_get_ ## name(VALUE self) { \
+  ruby_whisper_context_params *rwcp; \
+  GetContextParams(self, rwcp); \
+  return INT2NUM(rwcp->params.name); \
+} \
+static VALUE \
+ruby_whisper_context_params_set_ ## name(VALUE self, VALUE value) { \
+  ruby_whisper_context_params *rwcp; \
+  GetContextParams(self, rwcp); \
+  rwcp->params.name = NUM2INT(value); \
+  return value; \
+}
+
+#define DEFINE_PARAM(param_name, nth) \
+  id_ ## param_name = rb_intern(#param_name); \
+  param_names[nth] = id_ ## param_name; \
+  rb_define_method(cContextParams, #param_name, ruby_whisper_context_params_get_ ## param_name, 0); \
+  rb_define_method(cContextParams, #param_name "=", ruby_whisper_context_params_set_ ## param_name, 1);
+
+VALUE cContextParams;
+
+static ID param_names[NUM_PARAMS];
+static ID id_use_gpu;
+static ID id_flash_attn;
+static ID id_gpu_device;
+static ID id_dtw_token_timestamps;
+static ID id_dtw_aheads_preset;
+static ID id_dtw_n_top;
+
+static size_t
+ruby_whisper_context_params_memsize(const void *p)
+{
+  const ruby_whisper_context_params *rwcp = (ruby_whisper_context_params *)p;
+  if (!rwcp) {
+    return 0;
+  }
+  return sizeof(ruby_whisper_context_params);
+}
+
+const rb_data_type_t ruby_whisper_context_params_type = {
+  "ruby_whisper_context_params",
+  {0, RUBY_DEFAULT_FREE, ruby_whisper_context_params_memsize,},
+  0, 0,
+  0
+};
+
+static VALUE
+ruby_whisper_context_params_s_allocate(VALUE klass)
+{
+  ruby_whisper_context_params *rwcp;
+  return TypedData_Make_Struct(klass, ruby_whisper_context_params, &ruby_whisper_context_params_type, rwcp);
+}
+
+DEF_BOOLEAN_ATTR_METHOD(use_gpu);
+DEF_BOOLEAN_ATTR_METHOD(flash_attn);
+DEF_INT_ATTR_METHOD(gpu_device);
+DEF_BOOLEAN_ATTR_METHOD(dtw_token_timestamps);
+DEF_INT_ATTR_METHOD(dtw_aheads_preset);
+
+static VALUE
+ruby_whisper_context_params_get_dtw_n_top(VALUE self) {
+  ruby_whisper_context_params *rwcp;
+  GetContextParams(self, rwcp);
+
+  int dtw_n_top = rwcp->params.dtw_n_top;
+
+  return dtw_n_top == -1 ? Qnil : INT2NUM(dtw_n_top);
+}
+
+static VALUE
+ruby_whisper_context_params_set_dtw_n_top(VALUE self, VALUE value) {
+  ruby_whisper_context_params *rwcp;
+  GetContextParams(self, rwcp);
+
+  rwcp->params.dtw_n_top = NIL_P(value) ? -1 : NUM2INT(value);
+
+  return value;
+}
+
+#define SET_PARAM_IF_SAME(param_name) \
+  if (id == id_ ## param_name) { \
+    ruby_whisper_context_params_set_ ## param_name(self, value); \
+    continue; \
+  }
+
+static VALUE
+ruby_whisper_context_params_initialize(int argc, VALUE *argv, VALUE self)
+{
+  ruby_whisper_context_params *rwcp;
+  TypedData_Get_Struct(self, ruby_whisper_context_params, &ruby_whisper_context_params_type, rwcp);
+  rwcp->params = whisper_context_default_params();
+
+  VALUE kw_hash;
+  rb_scan_args_kw(RB_SCAN_ARGS_KEYWORDS, argc, argv, ":", &kw_hash);
+  if (NIL_P(kw_hash)) {
+    return Qnil;
+  }
+
+  VALUE values[NUM_PARAMS] = {Qundef};
+  rb_get_kwargs(kw_hash, param_names, 0, NUM_PARAMS, values);
+
+  ID id;
+  VALUE value;
+  for (int i = 0; i < NUM_PARAMS; i++) {
+    id = param_names[i];
+    value = values[i];
+    if (value == Qundef) {
+      continue;
+    }
+    SET_PARAM_IF_SAME(use_gpu)
+    SET_PARAM_IF_SAME(flash_attn)
+    SET_PARAM_IF_SAME(gpu_device)
+    SET_PARAM_IF_SAME(dtw_token_timestamps)
+    SET_PARAM_IF_SAME(dtw_aheads_preset)
+    SET_PARAM_IF_SAME(dtw_n_top)
+  }
+
+  return Qnil;
+}
+
+#undef SET_PARAM_IF_SAME
+
+void
+init_ruby_whisper_context_params(VALUE *cContext)
+{
+  cContextParams = rb_define_class_under(*cContext, "Params", rb_cObject);
+
+  rb_define_alloc_func(cContextParams, ruby_whisper_context_params_s_allocate);
+  rb_define_method(cContextParams, "initialize", ruby_whisper_context_params_initialize, -1);
+
+  DEFINE_PARAM(use_gpu, 0)
+  DEFINE_PARAM(flash_attn, 1)
+  DEFINE_PARAM(gpu_device, 2)
+  DEFINE_PARAM(dtw_token_timestamps, 3)
+  DEFINE_PARAM(dtw_aheads_preset, 4)
+  DEFINE_PARAM(dtw_n_top, 5)
+}
+
+#undef DEFINE_PARAM
+#undef DEF_INT_ATTR_METHOD
+#undef DEF_BOOLEAN_ATTR_METHOD
+#undef NUM_PARAMS

data/ext/ruby_whisper_model.c

@@ -1,4 +1,3 @@
-#include <ruby.h>
 #include "ruby_whisper.h"
 
 extern const rb_data_type_t ruby_whisper_type;

data/ext/ruby_whisper_params.c

@@ -1,4 +1,3 @@
-#include <ruby.h>
 #include "ruby_whisper.h"
 
 #define BOOL_PARAMS_SETTER(self, prop, value) \

data/ext/ruby_whisper_segment.c

@@ -1,4 +1,3 @@
-#include <ruby.h>
 #include "ruby_whisper.h"
 
 #define N_KEY_NAMES 6

data/ext/ruby_whisper_token.c

@@ -1,4 +1,3 @@
-#include <ruby.h>
 #include "ruby_whisper.h"
 
 #define N_KEY_NAMES 11
@@ -25,12 +24,34 @@ ruby_whisper_token_memsize(const void *p)
   if (!rwt) {
     return 0;
   }
-  return sizeof(rwt);
+  size_t size = sizeof(*rwt);
+  if (rwt->token_data) {
+    size += sizeof(*rwt->token_data);
+  }
+  return size;
+}
+
+static void
+ruby_whisper_token_mark(void *p)
+{
+  ruby_whisper_token *rwt = (ruby_whisper_token *)p;
+  rb_gc_mark(rwt->text);
+}
+
+static void
+ruby_whisper_token_free(void *p)
+{
+  ruby_whisper_token *rwt = (ruby_whisper_token *)p;
+  if (rwt->token_data) {
+    xfree(rwt->token_data);
+    rwt->token_data = NULL;
+  }
+  xfree(rwt);
 }
 
 static const rb_data_type_t ruby_whisper_token_type = {
   "ruby_whisper_token",
-  {0, RUBY_DEFAULT_FREE, ruby_whisper_token_memsize,},
+  {ruby_whisper_token_mark, ruby_whisper_token_free, ruby_whisper_token_memsize,},
   0, 0,
   0
 };
@@ -41,19 +62,19 @@ ruby_whisper_token_allocate(VALUE klass)
   ruby_whisper_token *rwt;
   VALUE token = TypedData_Make_Struct(klass, ruby_whisper_token, &ruby_whisper_token_type, rwt);
   rwt->token_data = NULL;
-  rwt->text = NULL;
+  rwt->text = Qnil;
   return token;
 }
 
 VALUE
 ruby_whisper_token_s_init(struct whisper_context *context, int i_segment, int i_token)
 {
-  whisper_token_data token_data = whisper_full_get_token_data(context, i_segment, i_token);
   const VALUE token = ruby_whisper_token_allocate(cToken);
   ruby_whisper_token *rwt;
   TypedData_Get_Struct(token, ruby_whisper_token, &ruby_whisper_token_type, rwt);
-  rwt->token_data = &token_data;
-  rwt->text = whisper_full_get_token_text(context, i_segment, i_token);
+  rwt->token_data = ALLOC(whisper_token_data);
+  *(rwt->token_data) = whisper_full_get_token_data(context, i_segment, i_token);
+  rwt->text = rb_str_new2(whisper_full_get_token_text(context, i_segment, i_token));
   return token;
 }
 
@@ -183,10 +204,9 @@ ruby_whisper_token_get_text(VALUE self)
 {
   ruby_whisper_token *rwt;
   GetToken(self, rwt);
-  return rb_str_new2(rwt->text);
+  return rwt->text;
 }
 
-
 /*
  * Start time of the token.
  *