Class: Cohere::Transcribe::ASR::NativeSession

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/asr/native.rb

Overview

One retained native Cohere model session.

Constant Summary collapse

MAX_NATIVE_SAMPLE_COUNT = 
2_147_483_647
NATIVE_CANCELLED_KIND = 
5
CANCELLATION_JOIN_INTERVAL = 
0.01
NATIVE_BATCH_STAT_FIELDS = 
%i[
  abi_version field_count total_us feature_wall_us feature_worker_us mel_pack_us
  encoder_graph_build_us encoder_graph_alloc_us encoder_input_us encoder_compute_us
  encoder_readback_us encoder_repack_us decoder_total_us decoder_cross_kv_us
  decoder_reserve_us decoder_decode_us render_us decoder_calls generation_steps
  encoder_microbatches token_id_readback_bytes
].freeze
NATIVE_BATCH_STAT_ABI_VERSION = 
1
NATIVE_FAILURE_KINDS = 
{
  1 => :fatal,
  2 => :oom,
  3 => :fatal,
  4 => :error
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(model_path, options, threads: nil, library: NativeLibrary.load) ⇒ NativeSession

Returns a new instance of NativeSession.



302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
 
# File 'lib/cohere/transcribe/asr/native.rb', line 302

def initialize(model_path, options, threads: nil, library: NativeLibrary.load)
  @library = library
  @model_path = Pathname.new(model_path).expand_path.freeze
  raise TranscriptionRuntimeError, "Native Dense model does not exist: #{@model_path}" unless @model_path.file?

  @mutex = Mutex.new
  @closed = false
  thread_count = normalize_threads(threads)
  @session = @library.open_session(
    model_path: @model_path,
    device: options.device.to_s,
    threads: thread_count
  )
  @default_max_new_tokens = Integer(options.max_new_tokens)
  @library.call(:session_set_max_new_tokens, @session, @default_max_new_tokens)
  @library.call(:session_set_beam_size, @session, 1)
  @library.call(
    :session_set_repetition_loop_guard,
    @session,
    options.stop_repetition_loops ? 1 : 0
  )
  @backend = @library.string(@library.call(:session_backend, @session)).freeze
  unless @backend == "cohere"
    close
    raise TranscriptionRuntimeError,
          "Native runtime selected #{@backend.inspect}, expected the Cohere Dense backend"
  end
  @compute_backend = @library.string(
    @library.call(:session_compute_backend, @session)
  ).freeze
  @device = canonical_device(@compute_backend)
  unless @device == options.device.to_s
    close
    raise TranscriptionRuntimeError,
          "Native runtime selected #{@compute_backend.inspect} (#{@device}) " \
          "after #{options.device.inspect} was resolved"
  end
  @batch_capacity = Integer(@library.call(:session_batch_capacity, @session))
  unless @batch_capacity.positive?
    close
    raise TranscriptionRuntimeError, "Native runtime reported an invalid Cohere batch capacity"
  end
rescue StandardError
  close if defined?(@session) && @session
  raise
end

Instance Attribute Details

#backend ⇒ Object (readonly)

Returns the value of attribute backend.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def backend
  @backend
end

#batch_capacity ⇒ Object (readonly)

Returns the value of attribute batch_capacity.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def batch_capacity
  @batch_capacity
end

#compute_backend ⇒ Object (readonly)

Returns the value of attribute compute_backend.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def compute_backend
  @compute_backend
end

#device ⇒ Object (readonly)

Returns the value of attribute device.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def device
  @device
end

#last_batch_metrics ⇒ Object (readonly)

Returns the value of attribute last_batch_metrics.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def last_batch_metrics
  @last_batch_metrics
end

#model_path ⇒ Object (readonly)

Returns the value of attribute model_path.



300
301
302
 
# File 'lib/cohere/transcribe/asr/native.rb', line 300

def model_path
  @model_path
end

Instance Method Details

#close ⇒ Object 



447
448
449
450
451
452
453
454
455
456
457
 
# File 'lib/cohere/transcribe/asr/native.rb', line 447

def close
  @mutex ||= Mutex.new
  @mutex.synchronize do
    return if @closed

    @library.call(:session_close, @session) if @library && @session
    @session = nil
    @closed = true
  end
  nil
end

#closed? ⇒ Boolean

Returns:

  • (Boolean) 


459
460
461
 
# File 'lib/cohere/transcribe/asr/native.rb', line 459

def closed?
  @closed
end

#memory ⇒ Object

Current free/total memory for the selected ggml device, used by the adaptive batch controller. Returns nil when a backend has no memory telemetry rather than guessing from process RSS.



466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
 
# File 'lib/cohere/transcribe/asr/native.rb', line 466

def memory
  @mutex.synchronize do
    ensure_open!
    free_bytes = [0].pack("Q")
    total_bytes = [0].pack("Q")
    status = @library.call(
      :session_memory,
      @session,
      Fiddle::Pointer[free_bytes],
      Fiddle::Pointer[total_bytes]
    )
    return nil unless status.zero?

    [free_bytes.unpack1("Q"), total_bytes.unpack1("Q")].freeze
  end
end

#transcribe(samples, language:, offset: 0.0, max_new_tokens: nil) ⇒ Object 



349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
 
# File 'lib/cohere/transcribe/asr/native.rb', line 349

def transcribe(samples, language:, offset: 0.0, max_new_tokens: nil)
  @mutex.synchronize do
    ensure_open!
    @last_batch_metrics = nil
    generation_limit = max_new_tokens.nil? ? @default_max_new_tokens : Integer(max_new_tokens)
    set_generation_limit!(generation_limit)
    binary, sample_count = float_samples(samples)
    raise TranscriptionRuntimeError, "Audio segment is too large for the native ABI" \
      if sample_count > MAX_NATIVE_SAMPLE_COUNT

    run_native_inference do
      result = @library.call(
        :session_transcribe_lang,
        @session,
        Fiddle::Pointer[binary],
        sample_count,
        Fiddle::Pointer["#{language}\0"]
      )
      raise_native_failure!("Native Cohere inference returned no result") if @library.null_pointer?(result)

      begin
        materialize_result(result, offset: Float(offset))
      ensure
        @library.call(:session_result_free, result)
      end
    end
  end
end

#transcribe_batch(sample_batches, language:, offsets: nil, max_new_tokens: nil) ⇒ Object

True padded-encoder/ragged-decoder batching in the native Cohere runtime. The capacity is queried from the loaded adapter because its logical decoder batch may span multiple encoder microbatches.



381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
 
# File 'lib/cohere/transcribe/asr/native.rb', line 381

def transcribe_batch(sample_batches, language:, offsets: nil, max_new_tokens: nil)
  unless sample_batches.is_a?(Array) && sample_batches.length.between?(1, @batch_capacity)
    raise ArgumentError,
          "sample_batches must contain between 1 and #{@batch_capacity} audio rows"
  end

  offsets ||= Array.new(sample_batches.length, 0.0)
  unless offsets.is_a?(Array) && offsets.length == sample_batches.length
    raise ArgumentError, "offsets must contain one value per audio row"
  end

  @mutex.synchronize do
    ensure_open!
    @last_batch_metrics = nil
    generation_limit = max_new_tokens.nil? ? @default_max_new_tokens : Integer(max_new_tokens)
    set_generation_limit!(generation_limit)

    buffers_and_counts = sample_batches.map { |samples| float_samples(samples) }
    oversized_lane = buffers_and_counts.index { |_binary, count| count > MAX_NATIVE_SAMPLE_COUNT }
    if oversized_lane
      raise TranscriptionRuntimeError,
            "Audio batch row #{oversized_lane} is too large for the native ABI"
    end

    run_native_inference do
      pointers = buffers_and_counts.map { |binary, _count| Fiddle::Pointer[binary] }
      pointer_table = pointers.map(&:to_i).pack("J*")
      counts = buffers_and_counts.map(&:last).pack("i!*")
      batch_result = @library.call(
        :session_transcribe_batch_lang,
        @session,
        Fiddle::Pointer[pointer_table],
        Fiddle::Pointer[counts],
        sample_batches.length,
        Fiddle::Pointer["#{language}\0"]
      )
      raise_native_failure!("Native Cohere batch inference returned no result") \
        if @library.null_pointer?(batch_result)

      begin
        @last_batch_metrics = materialize_batch_metrics(batch_result)
        count = @library.call(:session_batch_result_count, batch_result)
        unless count == sample_batches.length
          raise ExecutionError.new(
            "Native Cohere batch returned #{count} rows for #{sample_batches.length} inputs",
            failure_kind: :fatal
          )
        end
        Array.new(count) do |lane|
          result = @library.call(:session_batch_result_at, batch_result, lane)
          if @library.null_pointer?(result)
            raise ExecutionError.new(
              "Native Cohere batch result row #{lane} is missing",
              failure_kind: :fatal
            )
          end

          materialize_result(result, offset: Float(offsets.fetch(lane)))
        end.freeze
      ensure
        @library.call(:session_batch_result_free, batch_result)
      end
    end
  end
end