Add tests for Params#new_segment_callback=

Author: KitaitiMakoto

bindings/ruby/ext/ruby_whisper.cpp

@@ -35,6 +35,7 @@ extern "C" {
 VALUE mWhisper;
 VALUE cContext;
 VALUE cParams;
+VALUE cSegment;
 
 static VALUE ruby_whisper_s_lang_max_id(VALUE self) {
   return INT2NUM(whisper_lang_max_id());
@@ -476,6 +477,7 @@ void Init_whisper() {
   mWhisper = rb_define_module("Whisper");
   cContext = rb_define_class_under(mWhisper, "Context", rb_cObject);
   cParams  = rb_define_class_under(mWhisper, "Params", rb_cObject);
+  cSegment  = rb_define_class_under(mWhisper, "Segment", rb_cObject);
 
   rb_define_singleton_method(mWhisper, "lang_max_id", ruby_whisper_s_lang_max_id, 0);
   rb_define_singleton_method(mWhisper, "lang_id", ruby_whisper_s_lang_id, 1);

bindings/ruby/tests/test_whisper.rb

@@ -127,6 +127,56 @@ def test_whisper
     }
   end
 
+  def test_new_segment_callback
+    whisper = Whisper::Context.new(File.join(TOPDIR, '..', '..', 'models', 'ggml-base.en.bin'))
+
+    @params.new_segment_callback = ->(context, state, n_new, user_data) {
+      assert_kind_of Integer, n_new
+      assert n_new > 0
+      assert_same whisper, context
+
+      n_segments = context.full_n_segments
+      n_new.times do |i|
+        i_segment = n_segments - 1 + i
+        start_time = context.full_get_segment_t0(i_segment) * 10
+        end_time = context.full_get_segment_t1(i_segment) * 10
+        text = context.full_get_segment_text(i_segment)
+
+        assert_kind_of Integer, start_time
+        assert start_time >= 0
+        assert_kind_of Integer, end_time
+        assert end_time > 0
+        assert_match /ask not what your country can do for you, ask what you can do for your country/, text if i_segment == 0
+      end
+    }
+
+    jfk = File.join(TOPDIR, '..', '..', 'samples', 'jfk.wav')
+    whisper.transcribe(jfk, @params)
+  end
+
+  def test_new_segment_callback_closure
+    whisper = Whisper::Context.new(File.join(TOPDIR, '..', '..', 'models', 'ggml-base.en.bin'))
+
+    search_word = "what"
+    @params.new_segment_callback = ->(context, state, n_new, user_data) {
+      n_segments = context.full_n_segments
+      n_new.times do |i|
+        i_segment = n_segments - 1 + i
+        text = context.full_get_segment_text(i_segment)
+        if text.include?(search_word)
+          t0 = context.full_get_segment_t0(i_segment)
+          t1 = context.full_get_segment_t1(i_segment)
+          raise "search word '#{search_word}' found at between #{t0} and #{t1}"
+        end
+      end
+    }
+
+    jfk = File.join(TOPDIR, '..', '..', 'samples', 'jfk.wav')
+    assert_raise RuntimeError do
+      whisper.transcribe(jfk, @params)
+    end
+  end
+
   sub_test_case "After transcription" do
     class << self
       attr_reader :whisper