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.



346
347
348
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
377
378
379
380
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
 
# File 'lib/cohere/transcribe/asr/native.rb', line 346

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
  @session_ownership = SessionOwnership.new(@library)
  ObjectSpace.define_finalizer(
    self,
    self.class.send(:finalizer_for, @session_ownership)
  )
  thread_count = normalize_threads(threads)
  # An asynchronous exception may become deliverable immediately after
  # the foreign open returns. Defer it until the pointer belongs to the
  # independent cleanup state, then let the constructor rescue close it.
  Thread.handle_interrupt(Exception => :never) do
    session = @library.open_session(
      model_path: @model_path,
      device: options.device.to_s,
      threads: thread_count
    )
    @session_ownership.install(session)
    @session = session
  end
  @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 Exception # rubocop:disable Lint/RescueException -- native ownership must roll back on Interrupt too
  begin
    close if defined?(@session_ownership)
  rescue Exception # rubocop:disable Lint/RescueException -- preserve the constructor failure
    nil
  end
  raise
end

Instance Attribute Details

#backend ⇒ Object (readonly)

Returns the value of attribute backend.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def backend
  @backend
end

#batch_capacity ⇒ Object (readonly)

Returns the value of attribute batch_capacity.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def batch_capacity
  @batch_capacity
end

#compute_backend ⇒ Object (readonly)

Returns the value of attribute compute_backend.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def compute_backend
  @compute_backend
end

#device ⇒ Object (readonly)

Returns the value of attribute device.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def device
  @device
end

#last_batch_metrics ⇒ Object (readonly)

Returns the value of attribute last_batch_metrics.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def last_batch_metrics
  @last_batch_metrics
end

#model_path ⇒ Object (readonly)

Returns the value of attribute model_path.



344
345
346
 
# File 'lib/cohere/transcribe/asr/native.rb', line 344

def model_path
  @model_path
end

Instance Method Details

#close ⇒ Object 



507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
 
# File 'lib/cohere/transcribe/asr/native.rb', line 507

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

    @session = nil
    @closed = true
    begin
      @session_ownership&.close
    ensure
      ObjectSpace.undefine_finalizer(self)
    end
  end
  nil
end

#closed? ⇒ Boolean

Returns:

  • (Boolean) 


523
524
525
 
# File 'lib/cohere/transcribe/asr/native.rb', line 523

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.



530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
 
# File 'lib/cohere/transcribe/asr/native.rb', line 530

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 



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
 
# File 'lib/cohere/transcribe/asr/native.rb', line 409

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.



441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
 
# File 'lib/cohere/transcribe/asr/native.rb', line 441

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